Merge pull request #24029 from rthbound/dont-call-each-when-calling-body-on-response

Dont call each when calling body on response to fix #23964

Fixes #23964

6 files changed, 70 insertions, 19 deletions

diff --git a/actionpack/lib/action_controller/metal/live.rb b/actionpack/lib/action_controller/metal/live.rb
index 6055fde4f7..5d395cd8bd 100644
--- a/actionpack/lib/action_controller/metal/live.rb
+++ b/actionpack/lib/action_controller/metal/live.rb
@@ -163,14 +163,6 @@ module ActionController
         end
       end
 
-      def each
-        @response.sending!
-        while str = @buf.pop
-          yield str
-        end
-        @response.sent!
-      end
-
       # Write a 'close' event to the buffer; the producer/writing thread
       # uses this to notify us that it's finished supplying content.
       #
@@ -210,6 +202,14 @@ module ActionController
       def call_on_error
         @error_callback.call
       end
+
+      private
+
+        def each_chunk(&block)
+          while str = @buf.pop
+            yield str
+          end
+        end
     end
 
     class Response < ActionDispatch::Response #:nodoc: all
diff --git a/actionpack/lib/action_controller/test_case.rb b/actionpack/lib/action_controller/test_case.rb
index ecd21f29ce..ed2edcbe06 100644
--- a/actionpack/lib/action_controller/test_case.rb
+++ b/actionpack/lib/action_controller/test_case.rb
@@ -554,6 +554,8 @@ module ActionController
         end
         @request.query_string = ''
 
+        @response.sent!
+
         @response
       end
 
diff --git a/actionpack/lib/action_dispatch/http/response.rb b/actionpack/lib/action_dispatch/http/response.rb
index fa4c54701a..1515d59df3 100644
--- a/actionpack/lib/action_dispatch/http/response.rb
+++ b/actionpack/lib/action_dispatch/http/response.rb
@@ -68,7 +68,13 @@ module ActionDispatch # :nodoc:
     alias_method :headers,  :header
 
     delegate :[], :[]=, :to => :@header
-    delegate :each, :to => :@stream
+
+    def each(&block)
+      sending!
+      x = @stream.each(&block)
+      sent!
+      x
+    end
 
     CONTENT_TYPE = "Content-Type".freeze
     SET_COOKIE   = "Set-Cookie".freeze
@@ -97,10 +103,10 @@ module ActionDispatch # :nodoc:
 
       def body
         @str_body ||= begin
-                        buf = ''
-                        each { |chunk| buf << chunk }
-                        buf
-                      end
+          buf = ''
+            each { |chunk| buf << chunk }
+          buf
+        end
       end
 
       def write(string)
@@ -112,10 +118,13 @@ module ActionDispatch # :nodoc:
       end
 
       def each(&block)
-        @response.sending!
-        x = @buf.each(&block)
-        @response.sent!
-        x
+        if @str_body
+          return enum_for(:each) unless block_given?
+
+          yield @str_body
+        else
+          each_chunk(&block)
+        end
       end
 
       def abort
@@ -129,6 +138,12 @@ module ActionDispatch # :nodoc:
       def closed?
         @closed
       end
+
+      private
+
+        def each_chunk(&block)
+          @buf.each(&block) # extract into own method
+        end
     end
 
     def self.create(status = 200, header = {}, body = [], default_headers: self.default_headers)
diff --git a/actionpack/test/controller/live_stream_test.rb b/actionpack/test/controller/live_stream_test.rb
index a7759c080b..5977124594 100644
--- a/actionpack/test/controller/live_stream_test.rb
+++ b/actionpack/test/controller/live_stream_test.rb
@@ -246,7 +246,8 @@ module ActionController
 
     def assert_stream_closed
       assert response.stream.closed?, 'stream should be closed'
-      assert response.sent?, 'stream should be sent'
+      assert response.committed?,     'response should be committed'
+      assert response.sent?,          'response should be sent'
     end
 
     def capture_log_output
diff --git a/actionpack/test/dispatch/live_response_test.rb b/actionpack/test/dispatch/live_response_test.rb
index e4475f4233..de57c4be1d 100644
--- a/actionpack/test/dispatch/live_response_test.rb
+++ b/actionpack/test/dispatch/live_response_test.rb
@@ -65,7 +65,7 @@ module ActionController
         latch = Concurrent::CountDownLatch.new
 
         t = Thread.new {
-          @response.stream.each do
+          @response.each do
             latch.count_down
           end
         }
diff --git a/actionpack/test/dispatch/response_test.rb b/actionpack/test/dispatch/response_test.rb
index 658e0d004b..aa90433505 100644
--- a/actionpack/test/dispatch/response_test.rb
+++ b/actionpack/test/dispatch/response_test.rb
@@ -37,6 +37,39 @@ class ResponseTest < ActiveSupport::TestCase
     assert_equal "closed stream", e.message
   end
 
+  def test_each_isnt_called_if_str_body_is_written
+    # Controller writes and reads response body
+    each_counter = 0
+    @response.body = Object.new.tap {|o| o.singleton_class.send(:define_method, :each) { |&block| each_counter += 1; block.call 'foo' } }
+    @response['X-Foo'] = @response.body
+
+    assert_equal 1, each_counter, "#each was not called once"
+
+    # Build response
+    status, headers, body = @response.to_a
+
+    assert_equal 200, status
+    assert_equal "foo", headers['X-Foo']
+    assert_equal "foo", body.each.to_a.join
+
+    # Show that #each was not called twice
+    assert_equal 1, each_counter, "#each was not called once"
+  end
+
+  def test_set_header_after_read_body_during_action
+    @response.body
+
+    # set header after the action reads back @response.body
+    @response['x-header'] = "Best of all possible worlds."
+
+    # the response can be built.
+    status, headers, body = @response.to_a
+    assert_equal 200, status
+    assert_equal "", body.body
+
+    assert_equal "Best of all possible worlds.", headers['x-header']
+  end
+
   def test_read_body_during_action
     @response.body = "Hello, World!"