1 files changed, 62 insertions, 43 deletions

diff --git a/actionpack/lib/action_dispatch/middleware/flash.rb b/actionpack/lib/action_dispatch/middleware/flash.rb
index e90f8b9ce6..c51dcd542a 100644
--- a/actionpack/lib/action_dispatch/middleware/flash.rb
+++ b/actionpack/lib/action_dispatch/middleware/flash.rb
@@ -1,15 +1,6 @@
 require 'active_support/core_ext/hash/keys'
 
 module ActionDispatch
-  class Request < Rack::Request
-    # Access the contents of the flash. Use <tt>flash["notice"]</tt> to
-    # read a notice you put there or <tt>flash["notice"] = "hello"</tt>
-    # to put a new one.
-    def flash
-      @env[Flash::KEY] ||= Flash::FlashHash.from_session_value(session["flash"])
-    end
-  end
-
   # The flash provides a way to pass temporary primitive-types (String, Array, Hash) between actions. Anything you place in the flash will be exposed
   # to the very next action and then cleared out. This is a great way of doing notices and alerts, such as a create
   # action that sets <tt>flash[:notice] = "Post successfully created"</tt> before redirecting to a display action that can
@@ -47,6 +38,40 @@ module ActionDispatch
   class Flash
     KEY = 'action_dispatch.request.flash_hash'.freeze
 
+    module RequestMethods
+      # Access the contents of the flash. Use <tt>flash["notice"]</tt> to
+      # read a notice you put there or <tt>flash["notice"] = "hello"</tt>
+      # to put a new one.
+      def flash
+        flash = flash_hash
+        return flash if flash
+        self.flash = Flash::FlashHash.from_session_value(session["flash"])
+      end
+
+      def flash=(flash)
+        set_header Flash::KEY, flash
+      end
+
+      def flash_hash # :nodoc:
+        get_header Flash::KEY
+      end
+
+      def commit_flash # :nodoc:
+        session    = self.session || {}
+        flash_hash = self.flash_hash
+
+        if flash_hash && (flash_hash.present? || session.key?('flash'))
+          session["flash"] = flash_hash.to_session_value
+          self.flash = flash_hash.dup
+        end
+
+        if (!session.respond_to?(:loaded?) || session.loaded?) && # (reset_session uses {}, which doesn't implement #loaded?)
+            session.key?('flash') && session['flash'].nil?
+          session.delete('flash')
+        end
+      end
+    end
+
     class FlashNow #:nodoc:
       attr_accessor :flash
 
@@ -79,22 +104,31 @@ module ActionDispatch
     class FlashHash
       include Enumerable
 
-      def self.from_session_value(value)
-        flash = case value
-                when FlashHash # Rails 3.1, 3.2
-                  new(value.instance_variable_get(:@flashes), value.instance_variable_get(:@used))
-                when Hash # Rails 4.0
-                  new(value['flashes'], value['discard'])
-                else
-                  new
-                end
-
-        flash.tap(&:sweep)
+      def self.from_session_value(value) #:nodoc:
+        case value
+        when FlashHash # Rails 3.1, 3.2
+          flashes = value.instance_variable_get(:@flashes)
+          if discard = value.instance_variable_get(:@used)
+            flashes.except!(*discard)
+          end
+          new(flashes, flashes.keys)
+        when Hash # Rails 4.0
+          flashes = value['flashes']
+          if discard = value['discard']
+            flashes.except!(*discard)
+          end
+          new(flashes, flashes.keys)
+        else
+          new
+        end
       end
 
-      def to_session_value
-        return nil if empty?
-        {'discard' => @discard.to_a, 'flashes' => @flashes}
+      # Builds a hash containing the flashes to keep for the next request.
+      # If there are none to keep, returns nil.
+      def to_session_value #:nodoc:
+        flashes_to_keep = @flashes.except(*@discard)
+        return nil if flashes_to_keep.empty?
+        {'flashes' => flashes_to_keep}
       end
 
       def initialize(flashes = {}, discard = []) #:nodoc:
@@ -132,7 +166,7 @@ module ActionDispatch
       end
 
       def key?(name)
-        @flashes.key? name
+        @flashes.key? name.to_s
       end
 
       def delete(key)
@@ -249,25 +283,10 @@ module ActionDispatch
       end
     end
 
-    def initialize(app)
-      @app = app
-    end
-
-    def call(env)
-      @app.call(env)
-    ensure
-      session    = Request::Session.find(env) || {}
-      flash_hash = env[KEY]
-
-      if flash_hash && (flash_hash.present? || session.key?('flash'))
-        session["flash"] = flash_hash.to_session_value
-        env[KEY] = flash_hash.dup
-      end
+    def self.new(app) app; end
+  end
 
-      if (!session.respond_to?(:loaded?) || session.loaded?) && # (reset_session uses {}, which doesn't implement #loaded?)
-        session.key?('flash') && session['flash'].nil?
-        session.delete('flash')
-      end
-    end
+  class Request
+    prepend Flash::RequestMethods
   end
 end