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 @@ module ActionDispatch
     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 @@ module ActionDispatch
       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 @@ module ActionDispatch
       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 @@ module ActionDispatch
       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 @@ class FlashIntegrationTest < ActionDispatch::IntegrationTest
     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