From daa76a1e9021f811e73eb18ec10292e8ffd45aac Mon Sep 17 00:00:00 2001
From: Evan Phoenix <evan@phx.io>
Date: Tue, 5 Feb 2013 22:39:16 -0800
Subject: [PATCH] Add support for the rack hijack protocol

---
 lib/puma/client.rb  | 10 +++++++++-
 lib/puma/const.rb   |  4 ++++
 lib/puma/server.rb  | 30 ++++++++++++++++++++++++------
 test/hello-delay.ru |  3 +++
 test/hijack.ru      |  6 ++++++
 test/hijack2.ru     |  5 +++++
 6 files changed, 51 insertions(+), 7 deletions(-)
 create mode 100644 test/hello-delay.ru
 create mode 100644 test/hijack.ru
 create mode 100644 test/hijack2.ru

diff --git a/lib/puma/client.rb b/lib/puma/client.rb
index b47cfc46..ff22bf47 100644
--- a/lib/puma/client.rb
+++ b/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
diff --git a/lib/puma/const.rb b/lib/puma/const.rb
index b17ecedc..1c540b8d 100644
--- a/lib/puma/const.rb
+++ b/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
diff --git a/lib/puma/server.rb b/lib/puma/server.rb
index 9a68d21f..38cc269b 100644
--- a/lib/puma/server.rb
+++ b/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)
diff --git a/test/hello-delay.ru b/test/hello-delay.ru
new file mode 100644
index 00000000..85c899e2
--- /dev/null
+++ b/test/hello-delay.ru
@@ -0,0 +1,3 @@
+sleep 5
+
+run lambda { |env| [200, {"Content-Type" => "text/plain"}, ["Hello World"]] }
diff --git a/test/hijack.ru b/test/hijack.ru
new file mode 100644
index 00000000..b992f384
--- /dev/null
+++ b/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, {}, []]
+}
diff --git a/test/hijack2.ru b/test/hijack2.ru
new file mode 100644
index 00000000..bfdbeaf7
--- /dev/null
+++ b/test/hijack2.ru
@@ -0,0 +1,5 @@
+run lambda { |env|
+  body = lambda { |io| io.puts "BLAH\n"; io.close }
+
+  [200, { 'rack.hijack' => body }, []]
+}