From 6380f1a9f45e68f38480c0805cac62eb6708f72e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jos=C3=A9=20Valim?= <jose.valim@gmail.com>
Date: Tue, 19 Apr 2011 10:34:17 +0200
Subject: Be sure to not store the closed flash in the session.

---
 actionpack/lib/action_dispatch/middleware/flash.rb | 41 +++++++++++++---------
 actionpack/test/controller/flash_test.rb           |  9 ++++-
 2 files changed, 33 insertions(+), 17 deletions(-)

diff --git a/actionpack/lib/action_dispatch/middleware/flash.rb b/actionpack/lib/action_dispatch/middleware/flash.rb
index 027ff7f8ac..414405cc9e 100644
--- a/actionpack/lib/action_dispatch/middleware/flash.rb
+++ b/actionpack/lib/action_dispatch/middleware/flash.rb
@@ -4,7 +4,7 @@ module ActionDispatch
     # read a notice you put there or <tt>flash["notice"] = "hello"</tt>
     # to put a new one.
     def flash
-      @env['action_dispatch.request.flash_hash'] ||= (session["flash"] || Flash::FlashHash.new)
+      @env[Flash::KEY] ||= (session["flash"] || Flash::FlashHash.new)
     end
   end
 
@@ -40,18 +40,14 @@ module ActionDispatch
   #
   # See docs on the FlashHash class for more details about the flash.
   class Flash
+    KEY = 'action_dispatch.request.flash_hash'.freeze
+
     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, :flash if closed?
         @flash[k] = v
         @flash.discard(k)
         v
@@ -70,6 +66,10 @@ module ActionDispatch
       def notice=(message)
         self[:notice] = message
       end
+
+      def close!(new_flash)
+        @flash = new_flash
+      end
     end
 
     class FlashHash
@@ -81,10 +81,6 @@ module ActionDispatch
         @flashes = {}
       end
 
-      attr_reader :closed
-      alias :closed? :closed
-      def close!; @closed = true end
-
       def []=(k, v) #:nodoc:
         raise ClosedError, :flash if closed?
         keep(k)
@@ -152,6 +148,14 @@ module ActionDispatch
         @now ||= FlashNow.new(self)
       end
 
+      attr_reader :closed
+      alias :closed? :closed
+
+      def close!
+        @closed = true
+        @now.close!(self) if @now
+      end
+
       # Keeps either the entire current flash or a specific flash entry available for the next action:
       #
       #    flash.keep            # keeps the entire flash
@@ -231,13 +235,18 @@ module ActionDispatch
       @app.call(env)
     ensure
       session    = env['rack.session'] || {}
-      flash_hash = env['action_dispatch.request.flash_hash']
+      flash_hash = env[KEY]
 
       if flash_hash
-       if !flash_hash.empty? || session.key?('flash')
-        session["flash"] = flash_hash
-       end
-       flash_hash.close!
+        if !flash_hash.empty? || session.key?('flash')
+          session["flash"] = flash_hash
+          new_hash = flash_hash.dup
+        else
+          new_hash = flash_hash
+        end
+
+        env[KEY] = new_hash
+        new_hash.close!
       end
 
       if session.key?('flash') && session['flash'].empty?
diff --git a/actionpack/test/controller/flash_test.rb b/actionpack/test/controller/flash_test.rb
index 9c89f1334d..7b5bf8b21a 100644
--- a/actionpack/test/controller/flash_test.rb
+++ b/actionpack/test/controller/flash_test.rb
@@ -264,6 +264,14 @@ class FlashIntegrationTest < ActionDispatch::IntegrationTest
     end
   end
 
+  def test_setting_flash_does_not_raise_in_following_requests
+    with_test_route_set do
+      env = { 'action_dispatch.request.flash_hash' => ActionDispatch::Flash::FlashHash.new }
+      get '/set_flash', nil, env
+      get '/set_flash', nil, env
+    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 }
@@ -294,7 +302,6 @@ class FlashIntegrationTest < ActionDispatch::IntegrationTest
     end
   end
 
-
   private
 
     # Overwrite get to send SessionSecret in env hash
-- 
cgit v1.2.3