WebSocket support via partial hijack (#3058)

* Updates for response body processing

* Fixup resp_info[:no_body] to handle 101 'Switching Protocols'

* Create rack_conform.yml to test WebSocket

* rack_conform.yml - better sed regex, misc changes

* Create test_puma_server_partial_hijack.rb, pull 1 test from test_puma_server.rb

* rack_conform - use default branch, not websockets

* test_puma_server_partial_hijack.rb - minor refactor

* Revert "Fixup resp_info[:no_body] to handle 101 'Switching Protocols'"

This reverts commit 2abe219.

* request.rb - properly handle 101 'Switching Protocols' responses

* test_puma_server_partial_hijack.rb - TruffleRuby fix - add wait_readable

* test_puma_server_partial_hijack.rb - update comments re Truffle

* request.rb - change 'and' to '&&' in conditionals

* Remove `sleep 0.15 if Puma.jruby?` from test_puma_server_*.rb files

---------

Co-authored-by: MSP-Greg <Greg.mpls@gmail.com>

.github/workflows/rack_conform.yml

@@ -0,0 +1,67 @@
+name: rack-conform
+
+on: [push, pull_request, workflow_dispatch]
+
+permissions:
+  contents: read # to fetch code (actions/checkout)
+
+jobs:
+  skip_duplicate_runs:
+    uses: ./.github/workflows/skip_duplicate_workflow_runs.yaml
+
+  rack-conform:
+    name: >-
+      ${{ matrix.os }} Ruby ${{ matrix.ruby }} rack-conform
+    needs: skip_duplicate_runs
+    runs-on: ${{ matrix.os }}
+    if: |
+      !(   contains(github.event.pull_request.title,  '[ci skip]')
+        || contains(github.event.pull_request.title,  '[skip ci]')
+        || (needs.skip_duplicate_runs.outputs.should_skip == 'true'))
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - { os: ubuntu-20.04 , ruby: '2.7' }
+          - { os: ubuntu-22.04 , ruby: '3.2' }
+          - { os: ubuntu-22.04 , ruby: head  }
+
+    env:
+      BUNDLE_GEMFILE: gems/puma-head-rack-v3.rb
+      RACK_CONFORM_SERVER: puma
+      RACK_CONFORM_ENDPOINT: http://localhost:9292
+
+    steps:
+      - name: checkout rack-conform
+        uses: actions/checkout@v3
+        with:
+          repository: socketry/rack-conform
+
+      - name: Update gems/puma-head-rack-v3.rb
+        run: |
+          # use Puma from current repo (may be a fork) & sha
+          SRC="gem ['\"]puma['\"].*"
+          DST="gem 'puma', git: 'https://github.com/$GITHUB_REPOSITORY.git', ref: '$GITHUB_SHA'"
+          sed -i "s#$SRC#$DST#" gems/puma-head-rack-v3.rb
+
+      - name: load ruby
+        uses: ruby/setup-ruby@v1
+        with:
+          ruby-version: ${{ matrix.ruby }}
+          bundler-cache: true
+        timeout-minutes: 10
+
+      - name: cat gems/puma-head-rack-v3.rb.lock
+        run: cat gems/puma-head-rack-v3.rb.lock
+
+      - name: rack-conform test
+        id: test
+        timeout-minutes: 10
+        run: bundle exec bake test
+        continue-on-error: true
+        if: success()
+
+      - name: >-
+          Test outcome: ${{ steps.test.outcome }}
+        # every step must define a `uses` or `run` key
+        run: cat server.log

lib/puma/request.rb

@@ -173,7 +173,7 @@ def prepare_response(status, headers, res_body, requests, client)
       # 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)
+        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)
@@ -220,27 +220,33 @@ def prepare_response(status, headers, res_body, requests, client)
       cork_socket socket
 
       if resp_info[:no_body]
-        if content_length and status != 204
+        # 101 (Switching Protocols) doesn't return here or have content_length,
+        # it should be using `response_hijack`
+        unless status == 101
+          if content_length && status != 204
+            io_buffer.append CONTENT_LENGTH_S, content_length.to_s, line_ending
+          end
+
+          io_buffer << LINE_END
+          fast_write_str socket, io_buffer.read_and_reset
+          socket.flush
+          return keep_alive
+        end
+      else
+        if content_length
           io_buffer.append CONTENT_LENGTH_S, content_length.to_s, line_ending
+          chunked = false
+        elsif !response_hijack && resp_info[:allow_chunked]
+          io_buffer << TRANSFER_ENCODING_CHUNKED
+          chunked = true
         end
-
-        io_buffer << LINE_END
-        fast_write_str socket, io_buffer.read_and_reset
-        socket.flush
-        return keep_alive
-      end
-      if content_length
-        io_buffer.append CONTENT_LENGTH_S, content_length.to_s, line_ending
-        chunked = false
-      elsif !response_hijack and resp_info[:allow_chunked]
-        io_buffer << TRANSFER_ENCODING_CHUNKED
-        chunked = true
       end
 
       io_buffer << line_ending
 
       if response_hijack
         fast_write_str socket, io_buffer.read_and_reset
+        uncork_socket socket
         response_hijack.call socket
         return :async
       end
@@ -481,7 +487,7 @@ def req_env_post_parse(env)
       to_add = nil
 
       env.each do |k,v|
-        if k.start_with?("HTTP_") and k.include?(",") and k != "HTTP_TRANSFER,ENCODING"
+        if k.start_with?("HTTP_") && k.include?(",") && k != "HTTP_TRANSFER,ENCODING"
           if to_delete
             to_delete << k
           else

test/test_puma_server.rb

@@ -45,7 +45,6 @@ def server_run(**options, &block)
     @server = Puma::Server.new block || @app, @events, options
     @port = (@server.add_tcp_listener @host, 0).addr[1]
     @server.run
-    sleep 0.15 if Puma.jruby?
   end
 
   def header(sock)
@@ -689,23 +688,6 @@ def test_http_10_close_with_body
     assert_equal "HTTP/1.0 200 OK\r\nContent-Type: plain/text\r\nContent-Length: 5\r\n\r\nhello", data
   end
 
-  def test_http_10_partial_hijack_with_content_length
-    body_parts = ['abc', 'de']
-
-    server_run do
-      hijack_lambda = proc do | io |
-        io.write(body_parts[0])
-        io.write(body_parts[1])
-        io.close
-      end
-      [200, {"Content-Length" => "5", 'rack.hijack' => hijack_lambda}, nil]
-    end
-
-    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_http_10_keep_alive_without_body
     server_run { [204, {}, []] }
 

test/test_puma_server_partial_hijack.rb

@@ -0,0 +1,139 @@
+require_relative "helper"
+require "puma/events"
+require "puma/server"
+require "net/http"
+require "nio"
+require "ipaddr"
+
+class TestPumaServerPartialHijack < Minitest::Test
+  parallelize_me!
+
+  def setup
+    @host = "127.0.0.1"
+
+    @ios = []
+
+    @app = ->(env) { [200, {}, [env['rack.url_scheme']]] }
+
+    @log_writer = Puma::LogWriter.strings
+    @events = Puma::Events.new
+  end
+
+  def teardown
+    @server.stop(true)
+    assert_empty @log_writer.stdout.string
+    assert_empty @log_writer.stderr.string
+
+    # Errno::EBADF raised on macOS
+    @ios.each do |io|
+      begin
+        io.close if io.respond_to?(:close) && !io.closed?
+        File.unlink io.path if io.is_a? File
+      rescue Errno::EBADF
+      ensure
+        io = nil
+      end
+    end
+  end
+
+  def server_run(**options, &block)
+    options[:log_writer]  ||= @log_writer
+    options[:min_threads] ||= 1
+    @server = Puma::Server.new block || @app, @events, options
+    @port = (@server.add_tcp_listener @host, 0).addr[1]
+    @server.run
+  end
+
+  # only for shorter bodies!
+  def send_http_and_sysread(req)
+    send_http(req).sysread 2_048
+  end
+
+  def send_http_and_read(req)
+    send_http(req).read
+  end
+
+  def send_http(req)
+    new_connection << req
+  end
+
+  def new_connection
+    TCPSocket.new(@host, @port).tap {|sock| @ios << sock}
+  end
+
+  def test_101_body
+    headers = {
+      'Upgrade' => 'websocket',
+      'Connection' => 'Upgrade',
+      'Sec-WebSocket-Accept' => 's3pPLMBiTxaQ9kYGzzhZRbK+xOo=',
+      'Sec-WebSocket-Protocol' => 'chat'
+    }
+
+    body = -> (io) {
+      # below for TruffleRuby error with io.sysread
+      # Read Errno::EAGAIN: Resource temporarily unavailable
+      io.wait_readable 0.1
+      io.syswrite io.sysread(256)
+      io.close
+    }
+
+    server_run do |env|
+      [101, headers, body]
+    end
+
+    sock = new_connection
+    sock.syswrite "GET / HTTP/1.1\r\n\r\n"
+    resp = sock.sysread 1_024
+    echo_msg = "This should echo..."
+    sock.syswrite echo_msg
+
+    assert_includes resp, 'Connection: Upgrade'
+    assert_equal echo_msg, sock.sysread(256)
+  end
+
+  def test_101_header
+    headers = {
+      'Upgrade' => 'websocket',
+      'Connection' => 'Upgrade',
+      'Sec-WebSocket-Accept' => 's3pPLMBiTxaQ9kYGzzhZRbK+xOo=',
+      'Sec-WebSocket-Protocol' => 'chat',
+      'rack.hijack' => -> (io) {
+        # below for TruffleRuby error with io.sysread
+        # Read Errno::EAGAIN: Resource temporarily unavailable
+        io.wait_readable 0.1
+        io.syswrite io.sysread(256)
+        io.close
+      }
+    }
+
+    server_run do |env|
+      [101, headers, []]
+    end
+
+    sock = new_connection
+    sock.syswrite "GET / HTTP/1.1\r\n\r\n"
+    resp = sock.sysread 1_024
+    echo_msg = "This should echo..."
+    sock.syswrite echo_msg
+
+    assert_includes resp, 'Connection: Upgrade'
+    assert_equal echo_msg, sock.sysread(256)
+  end
+
+  def test_http_10_header_with_content_length
+    body_parts = ['abc', 'de']
+
+    server_run do
+      hijack_lambda = proc do | io |
+        io.write(body_parts[0])
+        io.write(body_parts[1])
+        io.close
+      end
+      [200, {"Content-Length" => "5", 'rack.hijack' => hijack_lambda}, nil]
+    end
+
+    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
+end