Add support for the rack hijack protocol

lib/puma/client.rb

@@ -36,14 +36,22 @@ module Puma
       @timeout_at = nil
 
       @requests_served = 0
+      @hijacked = false
     end
 
-    attr_reader :env, :to_io, :body, :io, :timeout_at, :ready
+    attr_reader :env, :to_io, :body, :io, :timeout_at, :ready, :hijacked
 
     def inspect
       "#<Puma::Client:0x#{object_id.to_s(16)} @ready=#{@ready.inspect}>"
     end
 
+    # For the hijack protocol (allows us to just put the Client object
+    # into the env)
+    def call
+      @hijacked = true
+      env[HIJACK_IO] ||= @io
+    end
+
     def set_timeout(val)
       @timeout_at = Time.now + val
     end

lib/puma/const.rb

@@ -147,5 +147,9 @@ module Puma
     COLON = ": ".freeze
 
     NEWLINE = "\n".freeze
+
+    HIJACK_P = "rack.hijack?".freeze
+    HIJACK = "rack.hijack".freeze
+    HIJACK_IO = "rack.hijack_io".freeze
   end
 end

lib/puma/server.rb

@@ -332,6 +332,9 @@ module Puma
 
       env[PUMA_SOCKET] = client
 
+      env[HIJACK_P] = true
+      env[HIJACK] = req
+
       body = req.body
 
       env[RACK_INPUT] = body
@@ -345,6 +348,9 @@ module Puma
       begin
         begin
           status, headers, res_body = @app.call(env)
+
+          return :async if req.hijacked
+
           status = status.to_i
 
           if status == -1
@@ -406,6 +412,8 @@ module Puma
           end
         end
 
+        response_hijack = nil
+
         headers.each do |k, vs|
           case k
           when CONTENT_LENGTH2
@@ -416,6 +424,9 @@ module Puma
             content_length = nil
           when CONTENT_TYPE
             next if no_body
+          when HIJACK
+            response_hijack = vs
+            next
           end
 
           vs.split(NEWLINE).each do |v|
@@ -435,18 +446,25 @@ module Puma
           lines << CONNECTION_CLOSE
         end
 
-        if content_length
-          lines.append CONTENT_LENGTH_S, content_length.to_s, line_ending
-          chunked = false
-        elsif allow_chunked
-          lines << TRANSFER_ENCODING_CHUNKED
-          chunked = true
+        unless response_hijack
+          if content_length
+            lines.append CONTENT_LENGTH_S, content_length.to_s, line_ending
+            chunked = false
+          elsif allow_chunked
+            lines << TRANSFER_ENCODING_CHUNKED
+            chunked = true
+          end
         end
 
         lines << line_ending
 
         fast_write client, lines.to_s
 
+        if response_hijack
+          response_hijack.call client
+          return :async
+        end
+
         res_body.each do |part|
           if chunked
             client.syswrite part.bytesize.to_s(16)

test/hello-delay.ru

@@ -0,0 +1,3 @@
+sleep 5
+
+run lambda { |env| [200, {"Content-Type" => "text/plain"}, ["Hello World"]] }

test/hijack.ru

@@ -0,0 +1,6 @@
+
+run lambda { |env|
+  io = env['rack.hijack'].call
+  io.puts "HTTP/1.1 200\r\n\r\nBLAH"
+  [-1, {}, []]
+}

test/hijack2.ru

@@ -0,0 +1,5 @@
+run lambda { |env|
+  body = lambda { |io| io.puts "BLAH\n"; io.close }
+
+  [200, { 'rack.hijack' => body }, []]
+}