hijacking support for Rack 1.5.x users

This requires Rack 1.5.x and unicorn 4.6.0 for hijacking
support.  Older versions of Rack continue to work fine,
but we must use unicorn 4.6.0 features to support this.

1 files changed, 34 insertions, 15 deletions

diff --git a/lib/rainbows/stream_response_epoll.rb b/lib/rainbows/stream_response_epoll.rb
index 3bb3540..33d7386 100644
--- a/lib/rainbows/stream_response_epoll.rb
+++ b/lib/rainbows/stream_response_epoll.rb
@@ -26,18 +26,24 @@ module Rainbows::StreamResponseEpoll
 
   def http_response_write(socket, status, headers, body)
     status = CODES[status.to_i] || status
-    ep_client = false
+    hijack = ep_client = false
 
     if headers
       # don't set extra headers here, this is only intended for
       # consuming by nginx.
       buf = "HTTP/1.0 #{status}\r\nStatus: #{status}\r\n"
       headers.each do |key, value|
-        if value =~ /\n/
-          # avoiding blank, key-only cookies with /\n+/
-          buf << value.split(/\n+/).map! { |v| "#{key}: #{v}\r\n" }.join
+        case key
+        when "rack.hijack"
+          hijack = hijack_prepare(value)
+          body = nil # ensure we do not close body
         else
-          buf << "#{key}: #{value}\r\n"
+          if /\n/ =~ value
+            # 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
       buf << HEADER_END
@@ -48,11 +54,22 @@ module Rainbows::StreamResponseEpoll
         buf = rv
       when :wait_writable
         ep_client = Client.new(socket, buf)
-        body.each { |chunk| ep_client.write(chunk) }
-        return ep_client.close
+        if hijack
+          ep_client.hijack(hijack)
+        else
+          body.each { |chunk| ep_client.write(chunk) }
+          ep_client.close
+        end
+        # body is nil on hijack, in which case ep_client is never closed by us
+        return
       end while true
     end
 
+    if hijack
+      hijack.call(socket)
+      return
+    end
+
     body.each do |chunk|
       if ep_client
         ep_client.write(chunk)
@@ -67,14 +84,15 @@ module Rainbows::StreamResponseEpoll
         end while true
       end
     end
-    ensure
-      body.respond_to?(:close) and body.close
-      if ep_client
-        ep_client.close
-      else
-        socket.shutdown
-        socket.close
-      end
+  ensure
+    return if hijack
+    body.respond_to?(:close) and body.close
+    if ep_client
+      ep_client.close
+    else
+      socket.shutdown
+      socket.close
+    end
   end
 
   # once a client is accepted, it is processed in its entirety here
@@ -88,6 +106,7 @@ module Rainbows::StreamResponseEpoll
       status, headers, body = @app.call(env)
     end
     @request.headers? or headers = nil
+    return if @request.hijacked?
     http_response_write(client, status, headers, body)
   rescue => e
     handle_error(client, e)