only write the jar if the response isn't committed

when streaming responses, we need to make sure the cookie jar is written
to the headers before returning up the stack. This commit introduces a
new method on the response object that writes the cookie jar to the
headers as the response is committed.  The middleware and test framework
will not write the cookie headers if the response has already been
committed.

fixes #14352

6 files changed, 57 insertions, 10 deletions

diff --git a/actionpack/lib/action_controller/metal/live.rb b/actionpack/lib/action_controller/metal/live.rb
index d60f1b0d44..43cf9b9723 100644
--- a/actionpack/lib/action_controller/metal/live.rb
+++ b/actionpack/lib/action_controller/metal/live.rb
@@ -177,13 +177,17 @@ module ActionController
         end
       end
 
-      def commit!
-        headers.freeze
+      private
+
+      def finalize_response
         super
+        jar = request.cookie_jar
+        # The response can be committed multiple times
+        jar.write self unless jar.committed?
+        jar.commit!
+        headers.freeze
       end
 
-      private
-
       def build_buffer(response, body)
         buf = Live::Buffer.new response
         body.each { |part| buf.write part }
diff --git a/actionpack/lib/action_controller/test_case.rb b/actionpack/lib/action_controller/test_case.rb
index 33a5858766..009f83861d 100644
--- a/actionpack/lib/action_controller/test_case.rb
+++ b/actionpack/lib/action_controller/test_case.rb
@@ -267,6 +267,18 @@ module ActionController
     def body
       @body ||= super
     end
+
+    # Was the response successful?
+    alias_method :success?, :successful?
+
+    # Was the URL not found?
+    alias_method :missing?, :not_found?
+
+    # Were we redirected?
+    alias_method :redirect?, :redirection?
+
+    # Was there a server-side error?
+    alias_method :error?, :server_error?
   end
 
   # Methods #destroy and #load! are overridden to avoid calling methods on the
@@ -583,7 +595,9 @@ module ActionController
         @controller.process(name)
 
         if cookies = @request.env['action_dispatch.cookies']
-          cookies.write(@response)
+          unless cookies.committed?
+            cookies.write(@response)
+          end
         end
         @response.prepare!
 
diff --git a/actionpack/lib/action_dispatch/http/response.rb b/actionpack/lib/action_dispatch/http/response.rb
index 2c6bcf7b7b..b5454f519f 100644
--- a/actionpack/lib/action_dispatch/http/response.rb
+++ b/actionpack/lib/action_dispatch/http/response.rb
@@ -140,6 +140,7 @@ module ActionDispatch # :nodoc:
 
     def commit!
       synchronize do
+        finalize_response
         @committed = true
         @cv.broadcast
       end
@@ -273,6 +274,9 @@ module ActionDispatch # :nodoc:
 
   private
 
+    def finalize_response
+    end
+
     def merge_default_headers(original, default)
       return original unless default.respond_to?(:merge)
 
diff --git a/actionpack/lib/action_dispatch/middleware/cookies.rb b/actionpack/lib/action_dispatch/middleware/cookies.rb
index 8b05cd6e11..c0039fa3f5 100644
--- a/actionpack/lib/action_dispatch/middleware/cookies.rb
+++ b/actionpack/lib/action_dispatch/middleware/cookies.rb
@@ -237,6 +237,15 @@ module ActionDispatch
         @secure = secure
         @options = options
         @cookies = {}
+        @committed = false
+      end
+
+      def committed?; @committed; end
+
+      def commit!
+        @committed = true
+        @set_cookies.freeze
+        @delete_cookies.freeze
       end
 
       def each(&block)
@@ -336,8 +345,8 @@ module ActionDispatch
       end
 
       def recycle! #:nodoc:
-        @set_cookies.clear
-        @delete_cookies.clear
+        @set_cookies = {}
+        @delete_cookies = {}
       end
 
       mattr_accessor :always_write_cookie
@@ -551,9 +560,11 @@ module ActionDispatch
       status, headers, body = @app.call(env)
 
       if cookie_jar = env['action_dispatch.cookies']
-        cookie_jar.write(headers)
-        if headers[HTTP_HEADER].respond_to?(:join)
-          headers[HTTP_HEADER] = headers[HTTP_HEADER].join("\n")
+        unless cookie_jar.committed?
+          cookie_jar.write(headers)
+          if headers[HTTP_HEADER].respond_to?(:join)
+            headers[HTTP_HEADER] = headers[HTTP_HEADER].join("\n")
+          end
         end
       end
 
diff --git a/actionpack/test/controller/live_stream_test.rb b/actionpack/test/controller/live_stream_test.rb
index eb8b2c9832..cc232f6b6e 100644
--- a/actionpack/test/controller/live_stream_test.rb
+++ b/actionpack/test/controller/live_stream_test.rb
@@ -102,6 +102,12 @@ module ActionController
         'test'
       end
 
+      def set_cookie
+        cookies[:hello] = "world"
+        response.stream.write "hello world"
+        response.close
+      end
+
       def render_text
         render :text => 'zomg'
       end
@@ -195,6 +201,13 @@ module ActionController
       end
     end
 
+    def test_set_cookie
+      @controller = TestController.new
+      get :set_cookie
+      assert_equal({'hello' => 'world'}, @response.cookies)
+      assert_equal "hello world", @response.body
+    end
+
     def test_set_response!
       @controller.set_response!(@request)
       assert_kind_of(Live::Response, @controller.response)
diff --git a/actionpack/test/dispatch/live_response_test.rb b/actionpack/test/dispatch/live_response_test.rb
index e0cfb73acf..336fdd569b 100644
--- a/actionpack/test/dispatch/live_response_test.rb
+++ b/actionpack/test/dispatch/live_response_test.rb
@@ -6,6 +6,7 @@ module ActionController
     class ResponseTest < ActiveSupport::TestCase
       def setup
         @response = Live::Response.new
+        @response.request = ActionDispatch::Request.new({}) #yolo
       end
 
       def test_header_merge