[PATCH] support for Rack::TempfileReaper middleware

[PATCH] support for Rack::TempfileReaper middleware

[Eric Wong]

Rack::TempfileReaper was added in rack 1.6 to cleanup temporary
files.  Make Yahns::TmpIO ducktype-compatible and put it into
env['rack.tempfiles'] array so Rack::TempfileReaper may be used to
free up space used by temporary buffer files.

ref: commit 3bdf5481e49d76b4502c51e5bdd93f68bfd1f0b4 in unicorn
---
 lib/yahns/http_client.rb  |  2 +-
 lib/yahns/http_context.rb | 19 +++++++++---------
 lib/yahns/tee_input.rb    |  2 +-
 lib/yahns/tmpio.rb        |  3 +++
 test/test_input.rb        | 51 ++++++++++++++++++++++++++++++++++++++++++++++-
 5 files changed, 65 insertions(+), 12 deletions(-)

diff --git a/lib/yahns/http_client.rb b/lib/yahns/http_client.rb
index a0fd5a4..620e925 100644
--- a/lib/yahns/http_client.rb
+++ b/lib/yahns/http_client.rb
@@ -57,7 +57,7 @@ class Yahns::HttpClient < Kgio::Socket # :nodoc:
             "Content-Length:#{len} too large (>#{mbs})", []
     end
     @state = :body
-    @input = k.tmpio_for(len)
+    @input = k.tmpio_for(len, @hs.env)
 
     rbuf = Thread.current[:yahns_rbuf]
     @hs.filter_body(rbuf, @hs.buf)
diff --git a/lib/yahns/http_context.rb b/lib/yahns/http_context.rb
index 1554086..8393ffe 100644
--- a/lib/yahns/http_context.rb
+++ b/lib/yahns/http_context.rb
@@ -77,14 +77,15 @@ module Yahns::HttpContext # :nodoc:
     @app_defaults["rack.errors"]
   end
 
-  def tmpio_for(len)
-    if len # Content-Length given
-      len <= @client_body_buffer_size ? StringIO.new("")
-           : Yahns::TmpIO.new(@input_buffer_tmpdir)
-    else # chunked, unknown length
-      mbs = @client_max_body_size
-      tmpdir = @input_buffer_tmpdir
-      mbs ? Yahns::CapInput.new(mbs, tmpdir) : Yahns::TmpIO.new(tmpdir)
-    end
+  def tmpio_for(len, env)
+    # short requests are most common
+    return StringIO.new('') if len && len <= @client_body_buffer_size;
+
+    # too big or chunked, unknown length
+    tmp = @input_buffer_tmpdir
+    mbs = @client_max_body_size
+    tmp = mbs ? Yahns::CapInput.new(mbs, tmp) : Yahns::TmpIO.new(tmp)
+    (env['rack.tempfiles'] ||= []) << tmp
+    tmp
   end
 end
diff --git a/lib/yahns/tee_input.rb b/lib/yahns/tee_input.rb
index 9028a6e..09933ca 100644
--- a/lib/yahns/tee_input.rb
+++ b/lib/yahns/tee_input.rb
@@ -19,7 +19,7 @@ class Yahns::TeeInput < Yahns::StreamInput # :nodoc:
   def initialize(client, request)
     @len = request.content_length
     super
-    @tmp = client.class.tmpio_for(@len)
+    @tmp = client.class.tmpio_for(@len, request.env)
   end
 
   # :call-seq:
diff --git a/lib/yahns/tmpio.rb b/lib/yahns/tmpio.rb
index ca86b4e..48832df 100644
--- a/lib/yahns/tmpio.rb
+++ b/lib/yahns/tmpio.rb
@@ -30,4 +30,7 @@ class Yahns::TmpIO < File # :nodoc:
     fp.sync = true
     fp
   end
+
+  # pretend we're Tempfile for Rack::TempfileReaper
+  alias close! close
 end
diff --git a/test/test_input.rb b/test/test_input.rb
index fe09a9a..63cf6ce 100644
--- a/test/test_input.rb
+++ b/test/test_input.rb
@@ -11,13 +11,28 @@ class TestInput < Testcase
 
   MD5 = lambda do |e|
     input = e["rack.input"]
+    tmp = e["rack.tempfiles"]
+    case input
+    when StringIO, Yahns::StreamInput
+      abort "unexpected tempfiles" if tmp && tmp.include?(input)
+    when Yahns::TmpIO
+      abort "rack.tempfiles missing" unless tmp
+      abort "rack.tempfiles missing rack.input" unless tmp.include?(input)
+    else
+      abort "unrecognized input type: #{input.class}"
+    end
+
     buf = ""
     md5 = Digest::MD5.new
     while input.read(16384, buf)
       md5 << buf
     end
     body = md5.hexdigest
-    h = { "Content-Length" => body.size.to_s, "Content-Type" => 'text/plain' }
+    h = {
+      "Content-Length" => body.size.to_s,
+      "Content-Type" => 'text/plain',
+      "X-Input-Class" => input.class.to_s,
+    }
     [ 200, h, [body] ]
   end
 
@@ -63,6 +78,40 @@ class TestInput < Testcase
     [ host, port, pid ]
   end
 
+  def test_big_buffer_true
+    host, port, pid = input_server(MD5, true)
+
+    c = get_tcp_client(host, port)
+    buf = 'hello'
+    c.write "PUT / HTTP/1.0\r\nContent-Length: 5\r\n\r\n#{buf}"
+    head, body = c.read.split(/\r\n\r\n/)
+    assert_match %r{^X-Input-Class: StringIO\r\n}, head
+    assert_equal Digest::MD5.hexdigest(buf), body
+    c.close
+
+    c = get_tcp_client(host, port)
+    buf = 'hello' * 10000
+    c.write "PUT / HTTP/1.0\r\nContent-Length: 50000\r\n\r\n#{buf}"
+    head, body = c.read.split(/\r\n\r\n/)
+
+    # TODO: shouldn't need CapInput with known Content-Length...
+    assert_match %r{^X-Input-Class: Yahns::(CapInput|TmpIO)\r\n}, head
+    assert_equal Digest::MD5.hexdigest(buf), body
+    c.close
+
+    c = get_tcp_client(host, port)
+    c.write "PUT / HTTP/1.1\r\nHost: example.com\r\nConnection: close\r\n"
+    c.write "Transfer-Encoding: chunked\r\n\r\n"
+    c.write "#{50000.to_s(16)}\r\n#{buf}\r\n0\r\n\r\n"
+    head, body = c.read.split(/\r\n\r\n/)
+    assert_match %r{^X-Input-Class: Yahns::CapInput\r\n}, head
+    assert_equal Digest::MD5.hexdigest(buf), body
+    c.close
+
+  ensure
+    quit_wait(pid)
+  end
+
   def test_read_negative_lazy; _read_neg(:lazy); end
   def test_read_negative_nobuffer; _read_neg(false); end
 
-- 
EW