raise if someone tries to modify the flash when it was already streamed back...

raise if someone tries to modify the flash when it was already streamed back to the client or converted to HTTP headers

actionpack/lib/action_dispatch/middleware/flash.rb

@@ -43,9 +43,15 @@ class Flash
     class FlashNow #:nodoc:
       def initialize(flash)
         @flash = flash
+        @closed = false
       end
 
+      attr_reader :closed
+      alias :closed? :closed
+      def close!; @closed = true end
+
       def []=(k, v)
+        raise ClosedError, "Cannot modify flash because it was closed. This means it was already streamed back to the client or converted to HTTP headers." if closed?
         @flash[k] = v
         @flash.discard(k)
         v
@@ -70,9 +76,15 @@ class FlashHash < Hash
       def initialize #:nodoc:
         super
         @used = Set.new
+        @closed = false
       end
 
+      attr_reader :closed
+      alias :closed? :closed
+      def close!; @closed = true end
+
       def []=(k, v) #:nodoc:
+        raise ClosedError, "Cannot modify flash because it was closed. This means it was already streamed back to the client or converted to HTTP headers." if closed?
         keep(k)
         super
       end
@@ -184,8 +196,11 @@ def call(env)
       session    = env['rack.session'] || {}
       flash_hash = env['action_dispatch.request.flash_hash']
 
-      if flash_hash && (!flash_hash.empty? || session.key?('flash'))
+      if flash_hash
+       if !flash_hash.empty? || session.key?('flash')
         session["flash"] = flash_hash
+       end
+       flash_hash.close!
       end
 
       if session.key?('flash') && session['flash'].empty?
@@ -193,4 +208,7 @@ def call(env)
       end
     end
   end
+
+  class ClosedError < StandardError #:nodoc:
+  end
 end

actionpack/test/controller/flash_test.rb

@@ -214,11 +214,20 @@ class FlashIntegrationTest < ActionDispatch::IntegrationTest
   SessionSecret = 'b3c631c314c0bbca50c1b2843150fe33'
 
   class TestController < ActionController::Base
+    def dont_set_flash
+      head :ok
+    end
+
     def set_flash
       flash["that"] = "hello"
       head :ok
     end
 
+    def set_flash_now
+      flash.now["that"] = "hello"
+      head :ok
+    end
+
     def use_flash
       render :inline => "flash: #{flash["that"]}"
     end
@@ -245,6 +254,47 @@ def test_just_using_flash_does_not_stream_a_cookie_back
     end
   end
 
+  def test_setting_flash_raises_after_stream_back_to_client
+    with_test_route_set do
+      env = { 'action_dispatch.request.flash_hash' => ActionDispatch::Flash::FlashHash.new }
+      get '/set_flash', nil, env
+      assert_raise(ActionDispatch::ClosedError) {
+        @request.flash['alert'] = 'alert'
+      }
+    end
+  end
+
+  def test_setting_flash_raises_after_stream_back_to_client_even_with_an_empty_flash
+    with_test_route_set do
+      env = { 'action_dispatch.request.flash_hash' => ActionDispatch::Flash::FlashHash.new }
+      get '/dont_set_flash', nil, env
+      assert_raise(ActionDispatch::ClosedError) {
+        @request.flash['alert'] = 'alert'
+      }
+    end
+  end
+
+  def test_setting_flash_now_raises_after_stream_back_to_client
+    with_test_route_set do
+      env = { 'action_dispatch.request.flash_hash' => ActionDispatch::Flash::FlashHash.new }
+      get '/set_flash_now', nil, env
+      assert_raise(ActionDispatch::ClosedError) {
+        @request.flash.now['alert'] = 'alert'
+      }
+    end
+  end
+
+  def test_setting_flash_now_raises_after_stream_back_to_client_even_with_an_empty_flash
+    with_test_route_set do
+      env = { 'action_dispatch.request.flash_hash' => ActionDispatch::Flash::FlashHash.new }
+      get '/dont_set_flash', nil, env
+      assert_raise(ActionDispatch::ClosedError) {
+        @request.flash.now['alert'] = 'alert'
+      }
+    end
+  end
+
+
   private
 
     # Overwrite get to send SessionSecret in env hash