From 02a072734906ac4c1ea77990207b84895ab4a7cb Mon Sep 17 00:00:00 2001
From: Eric Wong <e@80x24.org>
Date: Tue, 9 Jun 2015 20:17:18 +0000
Subject: ensure body is closed during hijack

Middlewares such as Rack::Lock (used by Rails) break badly unless
the response body is closed on hijack, so we will close it to follow
the lead of other popular Rack servers.

While it's unclear if there's anybody using rack.hijack with unicorn,
we'll try to emulate the behavior of other servers as much as
possible.

ref: https://github.com/ngauthier/tubesock/issues/10
---
 lib/unicorn/http_response.rb |  3 ---
 lib/unicorn/http_server.rb   | 20 +++++++++++++-------
 t/hijack.ru                  |  3 ++-
 t/t0200-rack-hijack.sh       |  7 +++++--
 test/unit/test_response.rb   | 11 -----------
 5 files changed, 20 insertions(+), 24 deletions(-)

diff --git a/lib/unicorn/http_response.rb b/lib/unicorn/http_response.rb
index 801bf9a..016dac8 100644
--- a/lib/unicorn/http_response.rb
+++ b/lib/unicorn/http_response.rb
@@ -52,12 +52,9 @@ module Unicorn::HttpResponse
     end
 
     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
 end
diff --git a/lib/unicorn/http_server.rb b/lib/unicorn/http_server.rb
index 2657c29..9129ed8 100644
--- a/lib/unicorn/http_server.rb
+++ b/lib/unicorn/http_server.rb
@@ -559,16 +559,22 @@ class Unicorn::HttpServer
   # in 3 easy steps: read request, call app, write app response
   def process_client(client)
     status, headers, body = @app.call(env = @request.read(client))
-    return if @request.hijacked?
 
-    if 100 == status.to_i
-      e100_response_write(client, env)
-      status, headers, body = @app.call(env)
+    begin
       return if @request.hijacked?
+
+      if 100 == status.to_i
+        e100_response_write(client, env)
+        status, headers, body = @app.call(env)
+        return if @request.hijacked?
+      end
+      @request.headers? or headers = nil
+      http_response_write(client, status, headers, body,
+                          @request.response_start_sent)
+    ensure
+      body.respond_to?(:close) and body.close
     end
-    @request.headers? or headers = nil
-    http_response_write(client, status, headers, body,
-                        @request.response_start_sent)
+
     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
diff --git a/t/hijack.ru b/t/hijack.ru
index fcb0b6d..4adec61 100644
--- a/t/hijack.ru
+++ b/t/hijack.ru
@@ -2,12 +2,13 @@ use Rack::Lint
 use Rack::ContentLength
 use Rack::ContentType, "text/plain"
 class DieIfUsed
+  @@n = 0
   def each
     abort "body.each called after response hijack\n"
   end
 
   def close
-    abort "body.close called after response hijack\n"
+    warn "closed DieIfUsed #{@@n += 1}\n"
   end
 end
 run lambda { |env|
diff --git a/t/t0200-rack-hijack.sh b/t/t0200-rack-hijack.sh
index f772071..de3eb82 100755
--- a/t/t0200-rack-hijack.sh
+++ b/t/t0200-rack-hijack.sh
@@ -16,12 +16,15 @@ t_begin "check response hijack" && {
 	test "xresponse.hijacked" = x"$(curl -sSfv http://$listen/hijack_res)"
 }
 
-t_begin "killing succeeds" && {
+t_begin "killing succeeds after hijack" && {
 	kill $unicorn_pid
 }
 
-t_begin "check stderr" && {
+t_begin "check stderr for hijacked body close" && {
 	check_stderr
+	grep 'closed DieIfUsed 1\>' $r_err
+	grep 'closed DieIfUsed 2\>' $r_err
+	! grep 'closed DieIfUsed 3\>' $r_err
 }
 
 t_done
diff --git a/test/unit/test_response.rb b/test/unit/test_response.rb
index d0f0c79..3478288 100644
--- a/test/unit/test_response.rb
+++ b/test/unit/test_response.rb
@@ -72,17 +72,6 @@ class ResponseTest < Test::Unit::TestCase
     assert ! out.closed?
   end
 
-  def test_body_closed
-    expect_body = %w(1 2 3 4).join("\n")
-    body = StringIO.new(expect_body)
-    body.rewind
-    out = StringIO.new
-    http_response_write(out,200, {}, body)
-    assert ! out.closed?
-    assert body.closed?
-    assert_match(expect_body, out.string.split(/\r\n/).last)
-  end
-
   def test_unknown_status_pass_through
     out = StringIO.new
     http_response_write(out,"666 I AM THE BEAST", {}, [] )
-- 
cgit v1.2.3-24-ge0c7