support for Rack hijack in request and response

Rack 1.5.0 (protocol version [1,2]) adds support for
hijacking the client socket (removing it from the control
of unicorn (or any other Rack webserver)).

Tested with rack 1.5.0.

3 files changed, 57 insertions, 10 deletions

diff --git a/lib/unicorn/http_request.rb b/lib/unicorn/http_request.rb
index 79ead2e..3bc64ed 100644
--- a/lib/unicorn/http_request.rb
+++ b/lib/unicorn/http_request.rb
@@ -91,6 +91,27 @@ class Unicorn::HttpParser
 
     e[RACK_INPUT] = 0 == content_length ?
                     NULL_IO : @@input_class.new(socket, self)
+    hijack_setup(e, socket)
     e.merge!(DEFAULTS)
   end
+
+  # Rack 1.5.0 (protocol version 1.2) adds hijack request support
+  if ((Rack::VERSION[0] << 8) | Rack::VERSION[1]) >= 0x0102
+    DEFAULTS["rack.hijack?"] = true
+
+    # FIXME: asking for clarification about this in
+    # http://mid.gmane.org/20130122100802.GA28585@dcvr.yhbt.net
+    DEFAULTS["rack.version"] = [1, 2]
+
+    RACK_HIJACK = "rack.hijack".freeze
+    RACK_HIJACK_IO = "rack.hijack_io".freeze
+
+    def hijack_setup(e, socket)
+      e[RACK_HIJACK] = proc { e[RACK_HIJACK_IO] ||= socket }
+    end
+  else
+    # old Rack, do nothing.
+    def hijack_setup(e, _)
+    end
+  end
 end
diff --git a/lib/unicorn/http_response.rb b/lib/unicorn/http_response.rb
index 579d957..083951c 100644
--- a/lib/unicorn/http_response.rb
+++ b/lib/unicorn/http_response.rb
@@ -25,6 +25,7 @@ module Unicorn::HttpResponse
   def http_response_write(socket, status, headers, body,
                           response_start_sent=false)
     status = CODES[status.to_i] || status
+    hijack = nil
 
     http_response_start = response_start_sent ? '' : 'HTTP/1.1 '
     if headers
@@ -33,19 +34,42 @@ module Unicorn::HttpResponse
             "Status: #{status}\r\n" \
             "Connection: close\r\n"
       headers.each do |key, value|
-        next if %r{\A(?:Date\z|Connection\z)}i =~ key
-        if value =~ /\n/
-          # avoiding blank, key-only cookies with /\n+/
-          buf << value.split(/\n+/).map! { |v| "#{key}: #{v}\r\n" }.join
+        case key
+        when %r{\A(?:Date\z|Connection\z)}i
+          next
+        when "rack.hijack"
+          # this was an illegal key in Rack < 1.5, so it should be
+          # OK to silently discard it for those older versions
+          hijack = hijack_prepare(value)
         else
-          buf << "#{key}: #{value}\r\n"
+          if value =~ /\n/
+            # avoiding blank, key-only cookies with /\n+/
+            buf << value.split(/\n+/).map! { |v| "#{key}: #{v}\r\n" }.join
+          else
+            buf << "#{key}: #{value}\r\n"
+          end
         end
       end
       socket.write(buf << CRLF)
     end
 
-    body.each { |chunk| socket.write(chunk) }
-    ensure
-      body.respond_to?(:close) and body.close
+    if hijack
+      body = nil # ensure we do not close body
+      hijack.call(socket)
+    else
+      body.each { |chunk| socket.write(chunk) }
+    end
+  ensure
+    body.respond_to?(:close) and body.close
+  end
+
+  # Rack 1.5.0 (protocol version 1.2) adds response hijacking support
+  if ((Rack::VERSION[0] << 8) | Rack::VERSION[1]) >= 0x0102
+    def hijack_prepare(value)
+      value
+    end
+  else
+    def hijack_prepare(_)
+    end
   end
 end
diff --git a/lib/unicorn/http_server.rb b/lib/unicorn/http_server.rb
index aa98aeb..2d8e4e1 100644
--- a/lib/unicorn/http_server.rb
+++ b/lib/unicorn/http_server.rb
@@ -559,8 +559,10 @@ class Unicorn::HttpServer
     @request.headers? or headers = nil
     http_response_write(client, status, headers, body,
                         @request.response_start_sent)
-    client.shutdown # in case of fork() in Rack app
-    client.close # flush and uncork socket immediately, no keepalive
+    unless client.closed? # rack.hijack may've close this for us
+      client.shutdown # in case of fork() in Rack app
+      client.close # flush and uncork socket immediately, no keepalive
+    end
   rescue => e
     handle_error(client, e)
   end