use the body proxy to freeze headers

avoid freezing the headers until the web server has actually read data
from the body proxy.  Once the webserver has read data, then we should
throw an error if someone tries to set a header

5 files changed, 57 insertions, 12 deletions

diff --git a/actionpack/lib/action_controller/metal/live.rb b/actionpack/lib/action_controller/metal/live.rb
index fe63cf2b98..41b997a755 100644
--- a/actionpack/lib/action_controller/metal/live.rb
+++ b/actionpack/lib/action_controller/metal/live.rb
@@ -125,9 +125,11 @@ module ActionController
       end
 
       def each
+        @response.sending!
         while str = @buf.pop
           yield str
         end
+        @response.sent!
       end
 
       def close
@@ -179,12 +181,16 @@ module ActionController
 
       private
 
-      def finalize_response
+      def before_committed
         super
         jar = request.cookie_jar
         # The response can be committed multiple times
         jar.write self unless committed?
-        jar.commit!
+      end
+
+      def before_sending
+        super
+        request.cookie_jar.commit!
         headers.freeze
       end
 
diff --git a/actionpack/lib/action_controller/test_case.rb b/actionpack/lib/action_controller/test_case.rb
index 009f83861d..e9166d8747 100644
--- a/actionpack/lib/action_controller/test_case.rb
+++ b/actionpack/lib/action_controller/test_case.rb
@@ -595,7 +595,7 @@ module ActionController
         @controller.process(name)
 
         if cookies = @request.env['action_dispatch.cookies']
-          unless cookies.committed?
+          unless @response.committed?
             cookies.write(@response)
           end
         end
diff --git a/actionpack/lib/action_dispatch/http/response.rb b/actionpack/lib/action_dispatch/http/response.rb
index b5454f519f..3d27ff2b24 100644
--- a/actionpack/lib/action_dispatch/http/response.rb
+++ b/actionpack/lib/action_dispatch/http/response.rb
@@ -91,7 +91,10 @@ module ActionDispatch # :nodoc:
       end
 
       def each(&block)
-        @buf.each(&block)
+        @response.sending!
+        x = @buf.each(&block)
+        @response.sent!
+        x
       end
 
       def close
@@ -118,6 +121,8 @@ module ActionDispatch # :nodoc:
       @blank        = false
       @cv           = new_cond
       @committed    = false
+      @sending      = false
+      @sent         = false
       @content_type = nil
       @charset      = nil
 
@@ -138,18 +143,37 @@ module ActionDispatch # :nodoc:
       end
     end
 
+    def await_sent
+      synchronize { @cv.wait_until { @sent } }
+    end
+
     def commit!
       synchronize do
-        finalize_response
+        before_committed
         @committed = true
         @cv.broadcast
       end
     end
 
-    def committed?
-      @committed
+    def sending!
+      synchronize do
+        before_sending
+        @sending = true
+        @cv.broadcast
+      end
     end
 
+    def sent!
+      synchronize do
+        @sent = true
+        @cv.broadcast
+      end
+    end
+
+    def sending?;   synchronize { @sending };   end
+    def committed?; synchronize { @committed }; end
+    def sent?;      synchronize { @sent };      end
+
     # Sets the HTTP status code.
     def status=(status)
       @status = Rack::Utils.status_code(status)
@@ -274,7 +298,10 @@ module ActionDispatch # :nodoc:
 
   private
 
-    def finalize_response
+    def before_committed
+    end
+
+    def before_sending
     end
 
     def merge_default_headers(original, default)
diff --git a/actionpack/test/controller/live_stream_test.rb b/actionpack/test/controller/live_stream_test.rb
index cc232f6b6e..2fd5c930ba 100644
--- a/actionpack/test/controller/live_stream_test.rb
+++ b/actionpack/test/controller/live_stream_test.rb
@@ -44,7 +44,7 @@ module ActionController
     tests SSETestController
 
     def wait_for_response_stream_close
-      response.stream.await_close
+      response.body
     end
 
     def test_basic_sse
@@ -186,8 +186,8 @@ module ActionController
     tests TestController
 
     def assert_stream_closed
-      response.stream.await_close
       assert response.stream.closed?, 'stream should be closed'
+      assert response.sent?, 'stream should be sent'
     end
 
     def capture_log_output
@@ -229,6 +229,7 @@ module ActionController
       @controller.response = @response
 
       t = Thread.new(@response) { |resp|
+        resp.await_commit
         resp.stream.each do |part|
           assert_equal parts.shift, part
           ol = @controller.latch
@@ -268,6 +269,7 @@ module ActionController
       assert_raises(ActionView::MissingTemplate) do
         get :exception_in_view
       end
+      assert response.body
       assert_stream_closed
     end
 
diff --git a/actionpack/test/dispatch/live_response_test.rb b/actionpack/test/dispatch/live_response_test.rb
index 336fdd569b..512f3a8a7a 100644
--- a/actionpack/test/dispatch/live_response_test.rb
+++ b/actionpack/test/dispatch/live_response_test.rb
@@ -35,6 +35,7 @@ module ActionController
           @response.stream.close
         }
 
+        @response.await_commit
         @response.each do |part|
           assert_equal 'foo', part
           latch.release
@@ -59,21 +60,30 @@ module ActionController
         assert_nil @response.headers['Content-Length']
       end
 
-      def test_headers_cannot_be_written_after_write
+      def test_headers_cannot_be_written_after_webserver_reads
         @response.stream.write 'omg'
+        latch = ActiveSupport::Concurrency::Latch.new
 
+        t = Thread.new {
+          @response.stream.each do |chunk|
+            latch.release
+          end
+        }
+
+        latch.await
         assert @response.headers.frozen?
         e = assert_raises(ActionDispatch::IllegalStateError) do
           @response.headers['Content-Length'] = "zomg"
         end
 
         assert_equal 'header already sent', e.message
+        @response.stream.close
+        t.join
       end
 
       def test_headers_cannot_be_written_after_close
         @response.stream.close
 
-        assert @response.headers.frozen?
         e = assert_raises(ActionDispatch::IllegalStateError) do
           @response.headers['Content-Length'] = "zomg"
         end