Rack::TempfileReaper

  - should do nothing (i.e. not bomb out) without env[rack.tempfiles]
  - should close env[rack.tempfiles] when body is closed
  - should initialize env[rack.tempfiles] when not already present
  - should append env[rack.tempfiles] when already present

3 files changed, 86 insertions, 0 deletions

diff --git a/lib/rack.rb b/lib/rack.rb
index 57119df3..341514c5 100644
--- a/lib/rack.rb
+++ b/lib/rack.rb
@@ -53,6 +53,7 @@ module Rack
   autoload :ShowExceptions, "rack/showexceptions"
   autoload :ShowStatus, "rack/showstatus"
   autoload :Static, "rack/static"
+  autoload :TempfileReaper, "rack/tempfile_reaper"
   autoload :URLMap, "rack/urlmap"
   autoload :Utils, "rack/utils"
   autoload :Multipart, "rack/multipart"
diff --git a/lib/rack/tempfile_reaper.rb b/lib/rack/tempfile_reaper.rb
new file mode 100644
index 00000000..1500b06a
--- /dev/null
+++ b/lib/rack/tempfile_reaper.rb
@@ -0,0 +1,22 @@
+require 'rack/body_proxy'
+
+module Rack
+
+  # Middleware tracks and cleans Tempfiles created throughout a request (i.e. Rack::Multipart)
+  # Ideas/strategy based on posts by Eric Wong and Charles Oliver Nutter
+  # https://groups.google.com/forum/#!searchin/rack-devel/temp/rack-devel/brK8eh-MByw/sw61oJJCGRMJ
+  class TempfileReaper
+    def initialize(app)
+      @app = app
+    end
+
+    def call(env)
+      env['rack.tempfiles'] ||= []
+      status, headers, body = @app.call(env)
+      body_proxy = BodyProxy.new(body) do
+        env['rack.tempfiles'].each { |f| f.close! } unless env['rack.tempfiles'].nil?
+      end
+      [status, headers, body_proxy]
+    end
+  end
+end
diff --git a/test/spec_tempfile_reaper.rb b/test/spec_tempfile_reaper.rb
new file mode 100644
index 00000000..ac39d878
--- /dev/null
+++ b/test/spec_tempfile_reaper.rb
@@ -0,0 +1,63 @@
+require 'rack/tempfile_reaper'
+require 'rack/lint'
+require 'rack/mock'
+
+describe Rack::TempfileReaper do
+  class MockTempfile
+    attr_reader :closed
+
+    def initialize
+      @closed = false
+    end
+
+    def close!
+      @closed = true
+    end
+  end
+
+  before do
+    @env = Rack::MockRequest.env_for
+  end
+
+  def call(app)
+    Rack::Lint.new(Rack::TempfileReaper.new(app)).call(@env)
+  end
+
+  should 'do nothing (i.e. not bomb out) without env[rack.tempfiles]' do
+    app = lambda { |_| [200, {}, ['Hello, World!']] }
+    response = call(app)
+    response[2].close
+    response[0].should.equal(200)
+  end
+
+  should 'close env[rack.tempfiles] when body is closed' do
+    tempfile1, tempfile2 = MockTempfile.new, MockTempfile.new
+    @env['rack.tempfiles'] = [ tempfile1, tempfile2 ]
+    app = lambda { |_| [200, {}, ['Hello, World!']] }
+    call(app)[2].close
+    tempfile1.closed.should.equal true
+    tempfile2.closed.should.equal true
+  end
+
+  should 'initialize env[rack.tempfiles] when not already present' do
+    tempfile = MockTempfile.new
+    app = lambda do |env|
+      env['rack.tempfiles'] << tempfile
+      [200, {}, ['Hello, World!']]
+    end
+    call(app)[2].close
+    tempfile.closed.should.equal true
+  end
+
+  should 'append env[rack.tempfiles] when already present' do
+    tempfile1, tempfile2 = MockTempfile.new, MockTempfile.new
+    @env['rack.tempfiles'] = [ tempfile1 ]
+    app = lambda do |env|
+      env['rack.tempfiles'] << tempfile2
+      [200, {}, ['Hello, World!']]
+    end
+    call(app)[2].close
+    tempfile1.closed.should.equal true
+    tempfile2.closed.should.equal true
+  end
+end