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

1 files changed, 19 insertions, 1 deletions

diff --git a/actionpack/lib/action_dispatch/middleware/flash.rb b/actionpack/lib/action_dispatch/middleware/flash.rb
index 21aeeb217a..410d3f7127 100644
--- a/actionpack/lib/action_dispatch/middleware/flash.rb
+++ b/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