Fix content-length calcuation in Rack:Response#write (#2150)

mattbrictson · ioquatix · commit c8b977f6c3a0 · 2024-01-31T20:41:23.000+13:00
When `Rack::Response` is initialized with an Array, it incorrectly
increments its internal `@length` value and emitted content-length
header on every subsequent write. The more times `write` is called, the
more the error accumulates.

This commit fixes the accumulation bug, and fixes/adds specs to properly
test the scenario where `write` is used multiple times.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -2,6 +2,12 @@
 
 All notable changes to this project will be documented in this file. For info on how to format all future additions to this file please reference [Keep A Changelog](https://keepachangelog.com/en/1.0.0/).
 
+## Unreleased
+
+### Changed
+
+- Fix incorrect content-length header that was emitted when `Rack::Response#write` was used in some situations. ([#2150](https://github.com/rack/rack/pull/2150), [@mattbrictson])
+
 ## [3.0.8] - 2023-06-14
 
 - Fix some unused variable verbose warnings. ([#2084](https://github.com/rack/rack/pull/2084), [@jeremyevans], [@skipkayhil](https://github.com/skipkayhil))
diff --git a/lib/rack/response.rb b/lib/rack/response.rb
@@ -328,6 +328,8 @@ def buffered_body!
             @body.each do |part|
               @length += part.to_s.bytesize
             end
+
+            @buffered = true
           elsif @body.respond_to?(:each)
             # Turn the user supplied body into a buffered array:
             body = @body
diff --git a/test/spec_response.rb b/test/spec_response.rb
@@ -412,7 +412,7 @@ def object_with_each.each
     status.must_equal 404
   end
 
-  it "correctly updates content-type when writing when not initialized with body" do
+  it "correctly updates content-length when writing when initialized without body" do
     r = Rack::Response.new
     r.write('foo')
     r.write('bar')
@@ -423,20 +423,39 @@ def object_with_each.each
     header['content-length'].must_equal '9'
   end
 
-  it "correctly updates content-type when writing when initialized with body" do
+  it "correctly updates content-length when writing when initialized with Array body" do
+    r = Rack::Response.new(["foo"])
+    r.write('bar')
+    r.write('baz')
+    _, header, body = r.finish
+    str = "".dup; body.each { |part| str << part }
+    str.must_equal "foobarbaz"
+    header['content-length'].must_equal '9'
+  end
+
+  it "correctly updates content-length when writing when initialized with String body" do
+    r = Rack::Response.new("foo")
+    r.write('bar')
+    r.write('baz')
+    _, header, body = r.finish
+    str = "".dup; body.each { |part| str << part }
+    str.must_equal "foobarbaz"
+    header['content-length'].must_equal '9'
+  end
+
+  it "correctly updates content-length when writing when initialized with object body that responds to #each" do
     obj = Object.new
     def obj.each
       yield 'foo'
       yield 'bar'
     end
-    ["foobar", ["foo", "bar"], obj].each do
-      r = Rack::Response.new(["foo", "bar"])
-      r.write('baz')
-      _, header, body = r.finish
-      str = "".dup; body.each { |part| str << part }
-      str.must_equal "foobarbaz"
-      header['content-length'].must_equal '9'
-    end
+    r = Rack::Response.new(obj)
+    r.write('baz')
+    r.write('baz')
+    _, header, body = r.finish
+    str = "".dup; body.each { |part| str << part }
+    str.must_equal "foobarbazbaz"
+    header['content-length'].must_equal '12'
   end
 
   it "doesn't return invalid responses" do
