From b72a86f66c722d56a6d77ed1d2779ace6ad103ed Mon Sep 17 00:00:00 2001
From: Eric Wong <normalperson@yhbt.net>
Date: Wed, 5 Jan 2011 22:39:03 -0800
Subject: close client socket after closing response body

Response bodies may capture the block passed to each
and save it for body.close, so don't close the socket
before we have a chance to call body.close
---
 lib/unicorn/http_response.rb |  1 -
 lib/unicorn/http_server.rb   |  1 +
 t/t0018-write-on-close.sh    | 23 +++++++++++++++++++++++
 t/write-on-close.ru          | 11 +++++++++++
 test/unit/test_response.rb   | 18 +++++++++---------
 5 files changed, 44 insertions(+), 10 deletions(-)
 create mode 100755 t/t0018-write-on-close.sh
 create mode 100644 t/write-on-close.ru

diff --git a/lib/unicorn/http_response.rb b/lib/unicorn/http_response.rb
index 3a03cd6..62b3ee9 100644
--- a/lib/unicorn/http_response.rb
+++ b/lib/unicorn/http_response.rb
@@ -38,7 +38,6 @@ module Unicorn::HttpResponse
     end
 
     body.each { |chunk| socket.write(chunk) }
-    socket.close # flushes and uncorks the socket immediately
     ensure
       body.respond_to?(:close) and body.close
   end
diff --git a/lib/unicorn/http_server.rb b/lib/unicorn/http_server.rb
index e2a4db7..3a6e51e 100644
--- a/lib/unicorn/http_server.rb
+++ b/lib/unicorn/http_server.rb
@@ -538,6 +538,7 @@ class Unicorn::HttpServer
     end
     @request.headers? or headers = nil
     http_response_write(client, status, headers, body)
+    client.close # flush and uncork socket immediately, no keepalive
   rescue => e
     handle_error(client, e)
   end
diff --git a/t/t0018-write-on-close.sh b/t/t0018-write-on-close.sh
new file mode 100755
index 0000000..3afefea
--- /dev/null
+++ b/t/t0018-write-on-close.sh
@@ -0,0 +1,23 @@
+#!/bin/sh
+. ./test-lib.sh
+t_plan 4 "write-on-close tests for funky response-bodies"
+
+t_begin "setup and start" && {
+	unicorn_setup
+	unicorn -D -c $unicorn_config write-on-close.ru
+	unicorn_wait_start
+}
+
+t_begin "write-on-close response body succeeds" && {
+	test xGoodbye = x"$(curl -sSf http://$listen/)"
+}
+
+t_begin "killing succeeds" && {
+	kill $unicorn_pid
+}
+
+t_begin "check stderr" && {
+	check_stderr
+}
+
+t_done
diff --git a/t/write-on-close.ru b/t/write-on-close.ru
new file mode 100644
index 0000000..54a2f2e
--- /dev/null
+++ b/t/write-on-close.ru
@@ -0,0 +1,11 @@
+class WriteOnClose
+  def each(&block)
+    @callback = block
+  end
+
+  def close
+    @callback.call "7\r\nGoodbye\r\n0\r\n\r\n"
+  end
+end
+use Rack::ContentType, "text/plain"
+run(lambda { |_| [ 200, [%w(Transfer-Encoding chunked)], WriteOnClose.new ] })
diff --git a/test/unit/test_response.rb b/test/unit/test_response.rb
index e2986c6..c3d9baf 100644
--- a/test/unit/test_response.rb
+++ b/test/unit/test_response.rb
@@ -27,7 +27,7 @@ class ResponseTest < Test::Unit::TestCase
   def test_response_headers
     out = StringIO.new
     http_response_write(out, 200, {"X-Whatever" => "stuff"}, ["cool"])
-    assert out.closed?
+    assert ! out.closed?
 
     assert out.length > 0, "output didn't have data"
   end
@@ -35,7 +35,7 @@ class ResponseTest < Test::Unit::TestCase
   def test_response_string_status
     out = StringIO.new
     http_response_write(out,'200', {}, [])
-    assert out.closed?
+    assert ! out.closed?
     assert out.length > 0, "output didn't have data"
     assert_equal 1, out.string.split(/\r\n/).grep(/^Status: 200 OK/).size
   end
@@ -43,7 +43,7 @@ class ResponseTest < Test::Unit::TestCase
   def test_response_200
     io = StringIO.new
     http_response_write(io, 200, {}, [])
-    assert io.closed?
+    assert ! io.closed?
     assert io.length > 0, "output didn't have data"
   end
 
@@ -51,7 +51,7 @@ class ResponseTest < Test::Unit::TestCase
     code = 400
     io = StringIO.new
     http_response_write(io, code, {}, [])
-    assert io.closed?
+    assert ! io.closed?
     lines = io.string.split(/\r\n/)
     assert_match(/.* Bad Request$/, lines.first,
                  "wrong default reason phrase")
@@ -60,7 +60,7 @@ class ResponseTest < Test::Unit::TestCase
   def test_rack_multivalue_headers
     out = StringIO.new
     http_response_write(out,200, {"X-Whatever" => "stuff\nbleh"}, [])
-    assert out.closed?
+    assert ! out.closed?
     assert_match(/^X-Whatever: stuff\r\nX-Whatever: bleh\r\n/, out.string)
   end
 
@@ -69,7 +69,7 @@ class ResponseTest < Test::Unit::TestCase
   def test_status_header_added
     out = StringIO.new
     http_response_write(out,200, {"X-Whatever" => "stuff"}, [])
-    assert out.closed?
+    assert ! out.closed?
     assert_equal 1, out.string.split(/\r\n/).grep(/^Status: 200 OK/i).size
   end
 
@@ -80,7 +80,7 @@ class ResponseTest < Test::Unit::TestCase
     out = StringIO.new
     header_hash = {"X-Whatever" => "stuff", 'StaTus' => "666" }
     http_response_write(out,200, header_hash, [])
-    assert out.closed?
+    assert ! out.closed?
     assert_equal 1, out.string.split(/\r\n/).grep(/^Status: 200 OK/i).size
     assert_equal 1, out.string.split(/\r\n/).grep(/^Status:/i).size
   end
@@ -91,7 +91,7 @@ class ResponseTest < Test::Unit::TestCase
     body.rewind
     out = StringIO.new
     http_response_write(out,200, {}, body)
-    assert out.closed?
+    assert ! out.closed?
     assert body.closed?
     assert_match(expect_body, out.string.split(/\r\n/).last)
   end
@@ -99,7 +99,7 @@ class ResponseTest < Test::Unit::TestCase
   def test_unknown_status_pass_through
     out = StringIO.new
     http_response_write(out,"666 I AM THE BEAST", {}, [] )
-    assert out.closed?
+    assert ! out.closed?
     headers = out.string.split(/\r\n\r\n/).first.split(/\r\n/)
     assert %r{\AHTTP/\d\.\d 666 I AM THE BEAST\z}.match(headers[0])
     status = headers.grep(/\AStatus:/i).first
-- 
cgit v1.2.3-24-ge0c7