From 0c5aded0922f80bd1a31c7d2a3974469a18160a8 Mon Sep 17 00:00:00 2001
From: Santiago Pastorino <santiago@wyeworks.com>
Date: Wed, 6 Apr 2011 12:05:58 -0300
Subject: [PATCH] raise if someone tries to modify the cookies when it was
 already streamed back to the client or converted to HTTP headers

---
 .../lib/action_dispatch/middleware/cookies.rb | 12 +++++
 actionpack/test/dispatch/cookies_test.rb      | 51 +++++++++++++++++++
 2 files changed, 63 insertions(+)

diff --git a/actionpack/lib/action_dispatch/middleware/cookies.rb b/actionpack/lib/action_dispatch/middleware/cookies.rb
index 7ac608f0a8..67c4b83d45 100644
--- a/actionpack/lib/action_dispatch/middleware/cookies.rb
+++ b/actionpack/lib/action_dispatch/middleware/cookies.rb
@@ -115,10 +115,15 @@ module ActionDispatch
         @delete_cookies = {}
         @host = host
         @secure = secure
+        @closed = false
 
         super()
       end
 
+      attr_reader :closed
+      alias :closed? :closed
+      def close!; @closed = true end
+
       # Returns the value of the cookie by +name+, or +nil+ if no such cookie exists.
       def [](name)
         super(name.to_s)
@@ -145,6 +150,7 @@ module ActionDispatch
       # Sets the cookie named +name+. The second argument may be the very cookie
       # value, or a hash of options as documented above.
       def []=(key, options)
+        raise ClosedError, :cookies if closed?
         if options.is_a?(Hash)
           options.symbolize_keys!
           value = options[:value]
@@ -225,6 +231,7 @@ module ActionDispatch
       end
 
       def []=(key, options)
+        raise ClosedError, :cookies if closed?
         if options.is_a?(Hash)
           options.symbolize_keys!
         else
@@ -263,6 +270,7 @@ module ActionDispatch
       end
 
       def []=(key, options)
+        raise ClosedError, :cookies if closed?
         if options.is_a?(Hash)
           options.symbolize_keys!
           options[:value] = @verifier.generate(options[:value])
@@ -305,6 +313,7 @@ module ActionDispatch
     end
 
     def call(env)
+      cookie_jar = nil
       status, headers, body = @app.call(env)
 
       if cookie_jar = env['action_dispatch.cookies']
@@ -315,6 +324,9 @@ module ActionDispatch
       end
 
       [status, headers, body]
+    ensure
+      cookie_jar = ActionDispatch::Request.new(env).cookie_jar unless cookie_jar
+      cookie_jar.close!
     end
   end
 end
diff --git a/actionpack/test/dispatch/cookies_test.rb b/actionpack/test/dispatch/cookies_test.rb
index 39159fd629..0d374e1d8b 100644
--- a/actionpack/test/dispatch/cookies_test.rb
+++ b/actionpack/test/dispatch/cookies_test.rb
@@ -495,3 +495,54 @@ class CookiesTest < ActionController::TestCase
       end
     end
 end
+
+class CookiesIntegrationTest < ActionDispatch::IntegrationTest
+  class TestController < ActionController::Base
+    def dont_set_cookies
+      head :ok
+    end
+
+    def set_cookies
+      cookies["that"] = "hello"
+      head :ok
+    end
+  end
+
+  def test_setting_cookies_raises_after_stream_back_to_client
+    with_test_route_set do
+      env = {}
+      get '/set_cookies', nil, env
+      assert_raise(ActionDispatch::ClosedError) {
+        request.cookie_jar['alert'] = 'alert'
+        cookies['alert'] = 'alert'
+      }
+    end
+  end
+
+  def test_setting_cookies_raises_after_stream_back_to_client_even_with_an_empty_flash
+    with_test_route_set do
+      env = {}
+      get '/dont_set_cookies', nil, {}
+      assert_raise(ActionDispatch::ClosedError) {
+        request.cookie_jar['alert'] = 'alert'
+      }
+    end
+  end
+
+  private
+
+  def with_test_route_set
+    with_routing do |set|
+      set.draw do
+        match ':action', :to => CookiesIntegrationTest::TestController
+      end
+
+      @app = self.class.build_app(set) do |middleware|
+        middleware.use ActionDispatch::Cookies
+        middleware.delete "ActionDispatch::ShowExceptions"
+      end
+
+      yield
+    end
+  end
+end