request.rb - fix partial hijack precedence, refactor, add tests

lib/puma/request.rb

@@ -106,6 +106,7 @@ def handle_request(client, requests)
         # is called
         res_body = app_body
 
+        # full hijack, app called env['rack.hijack']
         return :async if client.hijacked
 
         status = status.to_i
@@ -169,54 +170,55 @@ def prepare_response(status, headers, res_body, requests, client)
       resp_info = str_headers(env, status, headers, res_body, io_buffer, force_keep_alive)
 
       close_body = false
+      response_hijack = nil
+      content_length = resp_info[:content_length]
+      keep_alive     = resp_info[:keep_alive]
 
-      # below converts app_body into body, dependent on app_body's characteristics, and
-      # resp_info[:content_length] will be set if it can be determined
-      if !resp_info[: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
-          resp_info[:content_length] = body.sum(&:bytesize)
-        elsif res_body.is_a?(File) && res_body.respond_to?(:size)
-          body = res_body
-          resp_info[:content_length] = body.size
-        elsif res_body.respond_to?(:to_path) && res_body.respond_to?(:each) &&
+      if res_body.respond_to?(:each) && !resp_info[:response_hijack]
+        # below converts app_body into body, dependent on app_body's characteristics, and
+        # content_length will be set if it can be determined
+        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
+            content_length = body.sum(&:bytesize)
+          elsif res_body.is_a?(File) && res_body.respond_to?(:size)
+            body = res_body
+            content_length = body.size
+          elsif res_body.respond_to?(:to_path) && File.readable?(fn = res_body.to_path)
+            body = File.open fn, 'rb'
+            content_length = body.size
+            close_body = true
+          else
+            body = res_body
+          end
+        elsif !res_body.is_a?(::File) && res_body.respond_to?(:to_path) &&
             File.readable?(fn = res_body.to_path)
           body = File.open fn, 'rb'
-          resp_info[:content_length] = body.size
+          content_length = body.size
           close_body = true
+        elsif !res_body.is_a?(::File) && res_body.respond_to?(:filename) &&
+            res_body.respond_to?(:bytesize) && File.readable?(fn = res_body.filename)
+          # Sprockets::Asset
+          content_length = res_body.bytesize unless content_length
+          if (body_str = res_body.to_hash[:source])
+            body = [body_str]
+          else                           # avoid each and use a File object
+            body = File.open fn, 'rb'
+            close_body = true
+          end
         else
           body = res_body
         end
-      elsif !res_body.is_a?(::File) && res_body.respond_to?(:to_path) && res_body.respond_to?(:each) &&
-          File.readable?(fn = res_body.to_path)
-        body = File.open fn, 'rb'
-        resp_info[:content_length] = body.size
-        close_body = true
-      elsif !res_body.is_a?(::File) && res_body.respond_to?(:filename) && res_body.respond_to?(:each) &&
-          res_body.respond_to?(:bytesize) && File.readable?(fn = res_body.filename)
-        # Sprockets::Asset
-        resp_info[:content_length] = res_body.bytesize unless resp_info[:content_length]
-        if res_body.to_hash[:source]   # use each to return @source
-          body = res_body
-        else                           # avoid each and use a File object
-          body = File.open fn, 'rb'
-          close_body = true
-        end
       else
-        body = res_body
+        # partial hijack, from Rack spec:
+        #   Servers must ignore the body part of the response tuple when the
+        #   rack.hijack response header is present.
+        response_hijack = resp_info[:response_hijack] || res_body
       end
 
       line_ending = LINE_END
 
-      content_length = resp_info[:content_length]
-      keep_alive     = resp_info[:keep_alive]
-
-      if res_body && !res_body.respond_to?(:each)
-        response_hijack = res_body
-      else
-        response_hijack = resp_info[:response_hijack]
-      end
-
       cork_socket socket
 
       if resp_info[:no_body]
@@ -244,6 +246,8 @@ def prepare_response(status, headers, res_body, requests, client)
 
       io_buffer << line_ending
 
+      # partial hijack, we write headers, then hand the socket to the app via
+      # response_hijack.call
       if response_hijack
         fast_write_str socket, io_buffer.read_and_reset
         uncork_socket socket
@@ -358,16 +362,15 @@ def fast_write_response(socket, body, io_buffer, chunked, content_length)
         fast_write_str(socket, io_buffer.read_and_reset) unless io_buffer.length.zero?
       else
         # for enum bodies
-        fast_write_str socket, io_buffer.read_and_reset
         if chunked
           body.each do |part|
             next if (byte_size = part.bytesize).zero?
-             fast_write_str socket, (byte_size.to_s(16) << LINE_END)
-             fast_write_str socket, part
-             fast_write_str socket, LINE_END
+            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
         else
+          fast_write_str socket, io_buffer.read_and_reset
           body.each do |part|
             next if part.bytesize.zero?
             fast_write_str socket, part

test/test_puma_server_partial_hijack.rb → test/test_puma_server_hijack.rb

@@ -3,9 +3,18 @@
 require "puma/server"
 require "net/http"
 require "nio"
-require "ipaddr"
 
-class TestPumaServerPartialHijack < Minitest::Test
+require "rack"
+require "rack/body_proxy"
+
+# Tests check both the proper passing of the socket to the app, and also calling
+# of `body.close` on the response body.  Rack spec is unclear as to whether
+# calling close is expected.
+#
+# The sleep statements may not be needed for local CI, but are needed
+# for use with GitHub Actions...
+
+class TestPumaServerHijack < Minitest::Test
   parallelize_me!
 
   def setup
@@ -54,13 +63,39 @@ def send_http_and_read(req)
   end
 
   def send_http(req)
-    new_connection << req
+    t = new_connection
+    t.syswrite req
+    t
   end
 
   def new_connection
     TCPSocket.new(@host, @port).tap {|sock| @ios << sock}
   end
 
+  # Full hijack does not return headers
+  def test_full_hijack_body_close
+    @body_closed = false
+    server_run do |env|
+      io = env['rack.hijack'].call
+      io.syswrite 'Server listening'
+      io.wait_readable 2
+      io.syswrite io.sysread(256)
+      body = ::Rack::BodyProxy.new([]) { @body_closed = true }
+      [200, {}, body]
+    end
+
+    sock = send_http "GET / HTTP/1.1\r\n\r\n"
+
+    sock.wait_readable 2
+    assert_equal "Server listening", sock.sysread(256)
+
+    sock.syswrite "this should echo"
+    assert_equal "this should echo", sock.sysread(256)
+    Thread.pass
+    sleep 0.001 # intermittent failure, may need to increase in CI
+    assert @body_closed, "Reponse body must be closed"
+  end
+
   def test_101_body
     headers = {
       'Upgrade' => 'websocket',
@@ -81,8 +116,7 @@ def test_101_body
       [101, headers, body]
     end
 
-    sock = new_connection
-    sock.syswrite "GET / HTTP/1.1\r\n\r\n"
+    sock = send_http "GET / HTTP/1.1\r\n\r\n"
     resp = sock.sysread 1_024
     echo_msg = "This should echo..."
     sock.syswrite echo_msg
@@ -110,8 +144,7 @@ def test_101_header
       [101, headers, []]
     end
 
-    sock = new_connection
-    sock.syswrite "GET / HTTP/1.1\r\n\r\n"
+    sock = send_http "GET / HTTP/1.1\r\n\r\n"
     resp = sock.sysread 1_024
     echo_msg = "This should echo..."
     sock.syswrite echo_msg
@@ -132,8 +165,56 @@ def test_http_10_header_with_content_length
       [200, {"Content-Length" => "5", 'rack.hijack' => hijack_lambda}, nil]
     end
 
+    # using sysread may only receive part of the response
     data = send_http_and_read "GET / HTTP/1.0\r\nConnection: close\r\n\r\n"
 
     assert_equal "HTTP/1.0 200 OK\r\nContent-Length: 5\r\n\r\nabcde", data
   end
+
+  def test_partial_hijack_body_closes_body
+    @available = true
+    hdrs = { 'Content-Type' => 'text/plain' }
+    body = ::Rack::BodyProxy.new(HIJACK_LAMBDA) { @available = true }
+    partial_hijack_closes_body(hdrs, body)
+  end
+
+  def test_partial_hijack_header_closes_body_correct_precedence
+    @available = true
+    incorrect_lambda = ->(io) {
+      io.syswrite 'incorrect body.call'
+      io.close
+    }
+    hdrs = { 'Content-Type' => 'text/plain', 'rack.hijack' => HIJACK_LAMBDA}
+    body = ::Rack::BodyProxy.new(incorrect_lambda) { @available = true }
+    partial_hijack_closes_body(hdrs, body)
+  end
+
+  HIJACK_LAMBDA = ->(io) {
+    io.syswrite 'hijacked'
+    io.close
+  }
+
+  def partial_hijack_closes_body(hdrs, body)
+    server_run do
+      if @available
+        @available = false
+        [200, hdrs, body]
+      else
+        [500, { 'Content-Type' => 'text/plain' }, ['incorrect']]
+      end
+    end
+
+    sock1 = send_http "GET / HTTP/1.1\r\n\r\n"
+    sleep (Puma::IS_WINDOWS || !Puma::IS_MRI ? 0.3 : 0.1)
+    resp1 = sock1.sysread 1_024
+
+    sleep 0.01 # time for close block to be called ?
+
+    sock2 = send_http "GET / HTTP/1.1\r\n\r\n"
+    sleep (Puma::IS_WINDOWS || !Puma::IS_MRI ? 0.3 : 0.1)
+    resp2 = sock2.sysread 1_024
+
+    assert_operator resp1, :end_with?, 'hijacked'
+    assert_operator resp2, :end_with?, 'hijacked'
+  end
 end

test/test_rack_server.rb

@@ -215,43 +215,6 @@ def test_rack_body_proxy_content_length
     assert_equal str_ary_bytes, content_length
   end
 
-  def test_hijack_body_close
-    available = true
-    @server.app = ->(env) {
-      if available
-        available = false
-        hijack_lambda = ->(io) {
-          io.syswrite 'hijacked'
-          io.close
-        }
-        [200, { 'Content-Type' => 'text/plain', 'rack.hijack' => hijack_lambda},
-          ::Rack::BodyProxy.new([]) { available = true }]
-      else
-        [500, { 'Content-Type' => 'text/plain' }, ['incorrect']]
-      end
-    }
-
-    @server.run
-
-    socket1 = TCPSocket.new HOST, @port
-    socket1.syswrite "GET / HTTP/1.1\r\n\r\n"
-    sleep (Puma::IS_WINDOWS || !Puma::IS_MRI ? 0.25 : 0.1)
-    resp1 = socket1.sysread 1_024
-
-    sleep 0.01 # time for close block to be called ?
-
-    socket2 = TCPSocket.new HOST, @port
-    socket2.syswrite "GET / HTTP/1.1\r\n\r\n"
-    sleep (Puma::IS_WINDOWS || !Puma::IS_MRI ? 0.25 : 0.1)
-    resp2 = socket2.sysread 1_024
-
-    assert_operator resp1, :end_with?, 'hijacked'
-    assert_operator resp2, :end_with?, 'hijacked'
-
-    socket1.close
-    socket2.close
-  end
-
   def test_common_logger
     log = StringIO.new