diff options
| -rw-r--r-- | actionpack/lib/action_dispatch/middleware/flash.rb | 20 | ||||
| -rw-r--r-- | actionpack/test/controller/flash_test.rb | 50 |
2 files changed, 69 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 diff --git a/actionpack/test/controller/flash_test.rb b/actionpack/test/controller/flash_test.rb index 3569a2f213..a9c4f34275 100644 --- a/actionpack/test/controller/flash_test.rb +++ b/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 |