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