From a756c9263d9c575bd6b0401304367a76b84fcf88 Mon Sep 17 00:00:00 2001
From: Collin Sauve <collinsauve@gmail.com>
Date: Thu, 30 Mar 2023 23:27:49 -0400
Subject: [PATCH] Always write io_buffer when in "enum bodies" branch. (#3113)

* Fix 3112 - always write io_buffer when in "enum bodies" branch

* test_puma_server.rb - add 3 '404 empty body' tests

* Update request.rb

---------

Co-authored-by: MSP-Greg <Greg.mpls@gmail.com>
---
 lib/puma/request.rb      | 13 ++++++++++---
 test/test_puma_server.rb | 21 +++++++++++++++++++++
 2 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/lib/puma/request.rb b/lib/puma/request.rb
index 6cd4de19d7..72d4cc3803 100644
--- a/lib/puma/request.rb
+++ b/lib/puma/request.rb
@@ -180,7 +180,7 @@ def prepare_response(status, headers, res_body, requests, client)
         if !content_length && !resp_info[:transfer_encoding] && status != 204
           if res_body.respond_to?(:to_ary) && (array_body = res_body.to_ary) &&
               array_body.is_a?(Array)
-            body = array_body
+            body = array_body.compact
             content_length = body.sum(&:bytesize)
           elsif res_body.is_a?(File) && res_body.respond_to?(:size)
             body = res_body
@@ -363,12 +363,19 @@ def fast_write_response(socket, body, io_buffer, chunked, content_length)
       else
         # for enum bodies
         if chunked
+          empty_body = true
           body.each do |part|
-            next if (byte_size = part.bytesize).zero?
+            next if part.nil? || (byte_size = part.bytesize).zero?
+            empty_body = false
             io_buffer.append byte_size.to_s(16), LINE_END, part, LINE_END
             fast_write_str socket, io_buffer.read_and_reset
           end
-          fast_write_str socket, CLOSE_CHUNKED
+          if empty_body
+            io_buffer << CLOSE_CHUNKED
+            fast_write_str socket, io_buffer.read_and_reset
+          else
+            fast_write_str socket, CLOSE_CHUNKED
+          end
         else
           fast_write_str socket, io_buffer.read_and_reset
           body.each do |part|
diff --git a/test/test_puma_server.rb b/test/test_puma_server.rb
index aba7ece84f..0fed8b0520 100644
--- a/test/test_puma_server.rb
+++ b/test/test_puma_server.rb
@@ -1545,4 +1545,25 @@ def test_streaming_enum_body_2
     assert_equal str * loops, resp_body
     assert_operator times.last - times.first, :>, 1.0
   end
+
+  def test_empty_body_array_content_length_0
+    server_run { |env| [404, {'Content-Length' => '0'}, []] }
+
+    resp = send_http_and_sysread "GET / HTTP/1.1\r\n\r\n"
+    assert_equal "HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\n\r\n", resp
+  end
+
+  def test_empty_body_array_no_content_length
+    server_run { |env| [404, {}, []] }
+
+    resp = send_http_and_sysread "GET / HTTP/1.1\r\n\r\n"
+    assert_equal "HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\n\r\n", resp
+  end
+
+  def test_empty_body_enum
+    server_run { |env| [404, {}, [].to_enum] }
+
+    resp = send_http_and_sysread "GET / HTTP/1.1\r\n\r\n"
+    assert_equal "HTTP/1.1 404 Not Found\r\nTransfer-Encoding: chunked\r\n\r\n0\r\n\r\n", resp
+  end
 end