Move Flash into middleware

8 files changed, 223 insertions, 183 deletions

diff --git a/actionpack/lib/action_controller/metal/compatibility.rb b/actionpack/lib/action_controller/metal/compatibility.rb
index 7180c0ceda..0e869e4e87 100644
--- a/actionpack/lib/action_controller/metal/compatibility.rb
+++ b/actionpack/lib/action_controller/metal/compatibility.rb
@@ -31,7 +31,7 @@ module ActionController
                                              @variables_added @request_origin @url
                                              @parent_controller @action_name
                                              @before_filter_chain_aborted @_headers @_params
-                                             @_flash @_response)
+                                             @_response)
 
       # Controls the resource action separator
       cattr_accessor :resource_action_separator
diff --git a/actionpack/lib/action_controller/metal/flash.rb b/actionpack/lib/action_controller/metal/flash.rb
index 682f90e23b..bd768b634e 100644
--- a/actionpack/lib/action_controller/metal/flash.rb
+++ b/actionpack/lib/action_controller/metal/flash.rb
@@ -1,187 +1,14 @@
 module ActionController #:nodoc:
-  # The flash provides a way to pass temporary objects 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] = "Successfully created"</tt> before redirecting to a display action that can
-  # then expose the flash to its template. Actually, that exposure is automatically done. Example:
-  #
-  #   class PostsController < ActionController::Base
-  #     def create
-  #       # save post
-  #       flash[:notice] = "Successfully created post"
-  #       redirect_to posts_path(@post)
-  #     end
-  #
-  #     def show
-  #       # doesn't need to assign the flash notice to the template, that's done automatically
-  #     end
-  #   end
-  #
-  #   show.html.erb
-  #     <% if flash[:notice] %>
-  #       <div class="notice"><%= flash[:notice] %></div>
-  #     <% end %>
-  #
-  # This example just places a string in the flash, but you can put any object in there. And of course, you can put as
-  # many as you like at a time too. Just remember: They'll be gone by the time the next action has been performed.
-  #
-  # See docs on the FlashHash class for more details about the flash.
   module Flash
     extend ActiveSupport::Concern
 
     included do
+      delegate :flash, :to => :request
+      delegate :alert, :notice, :to => "request.flash"
       helper_method :alert, :notice
     end
 
-    class FlashNow #:nodoc:
-      def initialize(flash)
-        @flash = flash
-      end
-
-      def []=(k, v)
-        @flash[k] = v
-        @flash.discard(k)
-        v
-      end
-
-      def [](k)
-        @flash[k]
-      end
-    end
-
-    class FlashHash < Hash
-      def initialize #:nodoc:
-        super
-        @used = Set.new
-      end
-
-      def []=(k, v) #:nodoc:
-        keep(k)
-        super
-      end
-
-      def update(h) #:nodoc:
-        h.keys.each { |k| keep(k) }
-        super
-      end
-
-      alias :merge! :update
-
-      def replace(h) #:nodoc:
-        @used = Set.new
-        super
-      end
-
-      # Sets a flash that will not be available to the next action, only to the current.
-      #
-      #     flash.now[:message] = "Hello current action"
-      #
-      # This method enables you to use the flash as a central messaging system in your app.
-      # When you need to pass an object to the next action, you use the standard flash assign (<tt>[]=</tt>).
-      # When you need to pass an object to the current action, you use <tt>now</tt>, and your object will
-      # vanish when the current action is done.
-      #
-      # Entries set via <tt>now</tt> are accessed the same way as standard entries: <tt>flash['my-key']</tt>.
-      def now
-        FlashNow.new(self)
-      end
-
-      # Keeps either the entire current flash or a specific flash entry available for the next action:
-      #
-      #    flash.keep            # keeps the entire flash
-      #    flash.keep(:notice)   # keeps only the "notice" entry, the rest of the flash is discarded
-      def keep(k = nil)
-        use(k, false)
-      end
-
-      # Marks the entire flash or a single flash entry to be discarded by the end of the current action:
-      #
-      #     flash.discard              # discard the entire flash at the end of the current action
-      #     flash.discard(:warning)    # discard only the "warning" entry at the end of the current action
-      def discard(k = nil)
-        use(k)
-      end
-
-      # Mark for removal entries that were kept, and delete unkept ones.
-      #
-      # This method is called automatically by filters, so you generally don't need to care about it.
-      def sweep #:nodoc:
-        keys.each do |k|
-          unless @used.include?(k)
-            @used << k
-          else
-            delete(k)
-            @used.delete(k)
-          end
-        end
-
-        # clean up after keys that could have been left over by calling reject! or shift on the flash
-        (@used - keys).each{ |k| @used.delete(k) }
-      end
-
-      def store(session)
-        return if self.empty?
-        session["flash"] = self
-      end
-
-      private
-        # Used internally by the <tt>keep</tt> and <tt>discard</tt> methods
-        #     use()               # marks the entire flash as used
-        #     use('msg')          # marks the "msg" entry as used
-        #     use(nil, false)     # marks the entire flash as unused (keeps it around for one more action)
-        #     use('msg', false)   # marks the "msg" entry as unused (keeps it around for one more action)
-        # Returns the single value for the key you asked to be marked (un)used or the FlashHash itself
-        # if no key is passed.
-        def use(key = nil, used = true)
-          Array(key || keys).each { |k| used ? @used << k : @used.delete(k) }
-          return key ? self[key] : self
-        end
-    end
-
-    # 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 #:doc:
-      unless @_flash
-        @_flash = session["flash"] || FlashHash.new
-        @_flash.sweep
-      end
-
-      @_flash
-    end
-
-    # Convenience accessor for flash[:alert]
-    def alert
-      flash[:alert]
-    end
-
-    # Convenience accessor for flash[:alert]=
-    def alert=(message)
-      flash[:alert] = message
-    end
-
-    # Convenience accessor for flash[:notice]
-    def notice
-      flash[:notice]
-    end
-
-    # Convenience accessor for flash[:notice]=
-    def notice=(message)
-      flash[:notice] = message
-    end
-
     protected
-      def process_action(method_name)
-        @_flash = nil
-        super
-        @_flash.store(session) if @_flash
-        @_flash = nil
-      end
-
-      def reset_session
-        super
-        @_flash = nil
-      end
-
       def redirect_to(options = {}, response_status_and_flash = {}) #:doc:
         if alert = response_status_and_flash.delete(:alert)
           flash[:alert] = alert
diff --git a/actionpack/lib/action_controller/test_case.rb b/actionpack/lib/action_controller/test_case.rb
index 7e6da0e61d..14557ca782 100644
--- a/actionpack/lib/action_controller/test_case.rb
+++ b/actionpack/lib/action_controller/test_case.rb
@@ -239,13 +239,15 @@ module ActionController
       @request.assign_parameters(@controller.class.name.underscore.sub(/_controller$/, ''), action.to_s, parameters)
 
       @request.session = ActionController::TestSession.new(session) unless session.nil?
-      @request.session["flash"] = ActionController::Flash::FlashHash.new.update(flash) if flash
+      @request.session["flash"] = @request.flash.update(flash || {})
+      @request.session["flash"].sweep
 
       @controller.request = @request
       @controller.params.merge!(parameters)
       build_request_uri(action, parameters)
       Base.class_eval { include Testing }
       @controller.process_with_new_base_test(@request, @response)
+      @request.session.delete('flash') if @request.session['flash'].blank?
       @response
     end
 
diff --git a/actionpack/lib/action_dispatch.rb b/actionpack/lib/action_dispatch.rb
index f6d693d6f0..71bb67d3ff 100644
--- a/actionpack/lib/action_dispatch.rb
+++ b/actionpack/lib/action_dispatch.rb
@@ -43,6 +43,7 @@ module ActionDispatch
   autoload_under 'middleware' do
     autoload :Callbacks
     autoload :Cascade
+    autoload :Flash
     autoload :Head
     autoload :ParamsParser
     autoload :Rescue
diff --git a/actionpack/lib/action_dispatch/http/request.rb b/actionpack/lib/action_dispatch/http/request.rb
index 3d0aab8a06..22a08ec10d 100755
--- a/actionpack/lib/action_dispatch/http/request.rb
+++ b/actionpack/lib/action_dispatch/http/request.rb
@@ -463,10 +463,6 @@ EOM
       @env['rack.session.options'] = options
     end
 
-    def flash
-      session['flash'] || {}
-    end
-
     # Returns the authorization header regardless of whether it was specified directly or through one of the
     # proxy alternatives.
     def authorization
diff --git a/actionpack/lib/action_dispatch/middleware/flash.rb b/actionpack/lib/action_dispatch/middleware/flash.rb
new file mode 100644
index 0000000000..99b36366d6
--- /dev/null
+++ b/actionpack/lib/action_dispatch/middleware/flash.rb
@@ -0,0 +1,174 @@
+module ActionDispatch
+  class 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
+      session['flash'] ||= Flash::FlashHash.new
+    end
+  end
+
+  # The flash provides a way to pass temporary objects 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] = "Successfully created"</tt> before redirecting to a display action that can
+  # then expose the flash to its template. Actually, that exposure is automatically done. Example:
+  #
+  #   class PostsController < ActionController::Base
+  #     def create
+  #       # save post
+  #       flash[:notice] = "Successfully created post"
+  #       redirect_to posts_path(@post)
+  #     end
+  #
+  #     def show
+  #       # doesn't need to assign the flash notice to the template, that's done automatically
+  #     end
+  #   end
+  #
+  #   show.html.erb
+  #     <% if flash[:notice] %>
+  #       <div class="notice"><%= flash[:notice] %></div>
+  #     <% end %>
+  #
+  # This example just places a string in the flash, but you can put any object in there. And of course, you can put as
+  # many as you like at a time too. Just remember: They'll be gone by the time the next action has been performed.
+  #
+  # See docs on the FlashHash class for more details about the flash.
+  class Flash
+    class FlashNow #:nodoc:
+      def initialize(flash)
+        @flash = flash
+      end
+
+      def []=(k, v)
+        @flash[k] = v
+        @flash.discard(k)
+        v
+      end
+
+      def [](k)
+        @flash[k]
+      end
+    end
+
+    class FlashHash < Hash
+      def initialize #:nodoc:
+        super
+        @used = Set.new
+      end
+
+      def []=(k, v) #:nodoc:
+        keep(k)
+        super
+      end
+
+      def update(h) #:nodoc:
+        h.keys.each { |k| keep(k) }
+        super
+      end
+
+      alias :merge! :update
+
+      def replace(h) #:nodoc:
+        @used = Set.new
+        super
+      end
+
+      # Sets a flash that will not be available to the next action, only to the current.
+      #
+      #     flash.now[:message] = "Hello current action"
+      #
+      # This method enables you to use the flash as a central messaging system in your app.
+      # When you need to pass an object to the next action, you use the standard flash assign (<tt>[]=</tt>).
+      # When you need to pass an object to the current action, you use <tt>now</tt>, and your object will
+      # vanish when the current action is done.
+      #
+      # Entries set via <tt>now</tt> are accessed the same way as standard entries: <tt>flash['my-key']</tt>.
+      def now
+        FlashNow.new(self)
+      end
+
+      # Keeps either the entire current flash or a specific flash entry available for the next action:
+      #
+      #    flash.keep            # keeps the entire flash
+      #    flash.keep(:notice)   # keeps only the "notice" entry, the rest of the flash is discarded
+      def keep(k = nil)
+        use(k, false)
+      end
+
+      # Marks the entire flash or a single flash entry to be discarded by the end of the current action:
+      #
+      #     flash.discard              # discard the entire flash at the end of the current action
+      #     flash.discard(:warning)    # discard only the "warning" entry at the end of the current action
+      def discard(k = nil)
+        use(k)
+      end
+
+      # Mark for removal entries that were kept, and delete unkept ones.
+      #
+      # This method is called automatically by filters, so you generally don't need to care about it.
+      def sweep #:nodoc:
+        keys.each do |k|
+          unless @used.include?(k)
+            @used << k
+          else
+            delete(k)
+            @used.delete(k)
+          end
+        end
+
+        # clean up after keys that could have been left over by calling reject! or shift on the flash
+        (@used - keys).each{ |k| @used.delete(k) }
+      end
+
+      # Convenience accessor for flash[:alert]
+      def alert
+        self[:alert]
+      end
+
+      # Convenience accessor for flash[:alert]=
+      def alert=(message)
+        self[:alert] = message
+      end
+
+      # Convenience accessor for flash[:notice]
+      def notice
+        self[:notice]
+      end
+
+      # Convenience accessor for flash[:notice]=
+      def notice=(message)
+        self[:notice] = message
+      end
+
+      private
+        # Used internally by the <tt>keep</tt> and <tt>discard</tt> methods
+        #     use()               # marks the entire flash as used
+        #     use('msg')          # marks the "msg" entry as used
+        #     use(nil, false)     # marks the entire flash as unused (keeps it around for one more action)
+        #     use('msg', false)   # marks the "msg" entry as unused (keeps it around for one more action)
+        # Returns the single value for the key you asked to be marked (un)used or the FlashHash itself
+        # if no key is passed.
+        def use(key = nil, used = true)
+          Array(key || keys).each { |k| used ? @used << k : @used.delete(k) }
+          return key ? self[key] : self
+        end
+    end
+
+    def initialize(app)
+      @app = app
+    end
+
+    def call(env)
+      if (session = env['rack.session']) && (flash = session['flash'])
+        flash.sweep
+      end
+
+      @app.call(env)
+    ensure
+      if (session = env['rack.session']) && (flash = session['flash']) && flash.empty?
+        session.delete('flash')
+      end
+    end
+  end
+end
diff --git a/actionpack/test/abstract_unit.rb b/actionpack/test/abstract_unit.rb
index f26542b31e..10913c0fdb 100644
--- a/actionpack/test/abstract_unit.rb
+++ b/actionpack/test/abstract_unit.rb
@@ -87,10 +87,12 @@ end
 
 class ActionController::IntegrationTest < ActiveSupport::TestCase
   def self.build_app(routes = nil)
+    ActionDispatch::Flash
     ActionDispatch::MiddlewareStack.new { |middleware|
       middleware.use "ActionDispatch::ShowExceptions"
       middleware.use "ActionDispatch::Callbacks"
       middleware.use "ActionDispatch::ParamsParser"
+      middleware.use "ActionDispatch::Flash"
       middleware.use "ActionDispatch::Head"
     }.build(routes || ActionController::Routing::Routes)
   end
diff --git a/actionpack/test/controller/flash_test.rb b/actionpack/test/controller/flash_test.rb
index a9b60386f1..85a2e7f44b 100644
--- a/actionpack/test/controller/flash_test.rb
+++ b/actionpack/test/controller/flash_test.rb
@@ -159,7 +159,7 @@ class FlashTest < ActionController::TestCase
   end
 
   def test_keep_and_discard_return_values
-    flash = ActionController::Flash::FlashHash.new
+    flash = ActionDispatch::Flash::FlashHash.new
     flash.update(:foo => :foo_indeed, :bar => :bar_indeed)
 
     assert_equal(:foo_indeed, flash.discard(:foo)) # valid key passed
@@ -187,4 +187,42 @@ class FlashTest < ActionController::TestCase
     get :redirect_with_other_flashes
     assert_equal "Horses!", @controller.send(:flash)[:joyride]
   end
-end
\ No newline at end of file
+end
+
+class FlashIntegrationTest < ActionController::IntegrationTest
+  SessionKey = '_myapp_session'
+  SessionSecret = 'b3c631c314c0bbca50c1b2843150fe33'
+
+  class TestController < ActionController::Base
+    def set_flash
+      flash["that"] = "hello"
+      head :ok
+    end
+
+    def use_flash
+      render :inline => "flash: #{flash["that"]}"
+    end
+  end
+
+  def test_flash
+    with_test_route_set do
+      get '/set_flash'
+      assert_response :success
+      assert_equal "hello", @request.flash["that"]
+
+      get '/use_flash'
+      assert_response :success
+      assert_equal "flash: hello", @response.body
+    end
+  end
+
+  private
+    def with_test_route_set
+      with_routing do |set|
+        set.draw do |map|
+          match ':action', :to => ActionDispatch::Session::CookieStore.new(TestController, :key => SessionKey, :secret => SessionSecret)
+        end
+        yield
+      end
+    end
+end