about summary refs log tree commit homepage
diff options
context:
space:
mode:
-rw-r--r--lib/unicorn/tee_input.rb9
-rw-r--r--lib/unicorn/tmpio.rb3
-rw-r--r--test/unit/test_tee_input.rb10
3 files changed, 21 insertions, 1 deletions
diff --git a/lib/unicorn/tee_input.rb b/lib/unicorn/tee_input.rb
index 637c583..3c6d18a 100644
--- a/lib/unicorn/tee_input.rb
+++ b/lib/unicorn/tee_input.rb
@@ -28,13 +28,20 @@ class Unicorn::TeeInput < Unicorn::StreamInput
     @@client_body_buffer_size
   end
 
+  # for Rack::TempfileReaper in rack 1.6+
+  def new_tmpio # :nodoc:
+    tmpio = Unicorn::TmpIO.new
+    (@parser.env['rack.tempfiles'] ||= []) << tmpio
+    tmpio
+  end
+
   # Initializes a new TeeInput object.  You normally do not have to call
   # this unless you are writing an HTTP server.
   def initialize(socket, request)
     @len = request.content_length
     super
     @tmp = @len && @len <= @@client_body_buffer_size ?
-           StringIO.new("") : Unicorn::TmpIO.new
+           StringIO.new("") : new_tmpio
   end
 
   # :call-seq:
diff --git a/lib/unicorn/tmpio.rb b/lib/unicorn/tmpio.rb
index 2da05a2..dcdf9da 100644
--- a/lib/unicorn/tmpio.rb
+++ b/lib/unicorn/tmpio.rb
@@ -26,4 +26,7 @@ class Unicorn::TmpIO < File
   def size
     stat.size
   end unless File.method_defined?(:size)
+
+  # pretend we're Tempfile for Rack::TempfileReaper
+  alias close! close
 end
diff --git a/test/unit/test_tee_input.rb b/test/unit/test_tee_input.rb
index 0c2c941..4647e66 100644
--- a/test/unit/test_tee_input.rb
+++ b/test/unit/test_tee_input.rb
@@ -29,6 +29,13 @@ class TestTeeInput < Test::Unit::TestCase
     end while true
   end
 
+  def check_tempfiles
+    tmp = @parser.env["rack.tempfiles"]
+    assert_instance_of Array, tmp
+    assert_operator tmp.size, :>=, 1
+    assert_instance_of Unicorn::TmpIO, tmp[0]
+  end
+
   def test_gets_long
     r = init_request("hello", 5 + (4096 * 4 * 3) + "#$/foo#$/".size)
     ti = TeeInput.new(@rd, r)
@@ -106,6 +113,7 @@ class TestTeeInput < Test::Unit::TestCase
     assert_kind_of File, ti.tmp
     assert_equal 0, ti.tmp.pos
     assert_equal Unicorn::Const::MAX_BODY + 1, ti.size
+    check_tempfiles
   end
 
   def test_read_in_full_if_content_length
@@ -148,6 +156,7 @@ class TestTeeInput < Test::Unit::TestCase
     assert_nil ti.read(1)
     pid, status = Process.waitpid2(pid)
     assert status.success?
+    check_tempfiles
   end
 
   def test_chunked
@@ -183,6 +192,7 @@ class TestTeeInput < Test::Unit::TestCase
     status = nil
     pid, status = Process.waitpid2(pid)
     assert status.success?
+    check_tempfiles
   end
 
   def test_chunked_ping_pong