Eliminate dependency on garbage collection for cleanup of Tempfiles created by Rack for multipart form data

lib/rack.rb

@@ -53,6 +53,7 @@ def self.release
   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"

lib/rack/multipart/parser.rb

@@ -16,10 +16,10 @@ def self.create(env)
         content_length = env['CONTENT_LENGTH']
         content_length = content_length.to_i if content_length
 
-        new($1, io, content_length)
+        new($1, io, content_length, env)
       end
 
-      def initialize(boundary, io, content_length)
+      def initialize(boundary, io, content_length, env)
         @buf            = ""
 
         if @buf.respond_to? :force_encoding
@@ -31,6 +31,7 @@ def initialize(boundary, io, content_length)
         @io             = io
         @content_length = content_length
         @boundary_size  = Utils.bytesize(@boundary) + EOL.size
+        @env = env
 
         if @content_length
           @content_length -= @boundary_size
@@ -112,7 +113,7 @@ def get_current_head_and_filename_and_content_type_and_name_and_body
             filename = get_filename(head)
 
             if filename
-              body = Tempfile.new("RackMultipart")
+              (@env['rack.tempfiles'] ||= []) << body = Tempfile.new("RackMultipart")
               body.binmode  if body.respond_to?(:binmode)
             end
 

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

test/spec_request.rb

@@ -740,6 +740,31 @@
     req.POST["mean"][:tempfile].read.should.equal "--AaB03xha"
   end
 
+  should "record tempfiles from multipart form data in env[rack.tempfiles]" do
+    input = <<EOF
+--AaB03x\r
+content-disposition: form-data; name="fileupload"; filename="foo.jpg"\r
+Content-Type: image/jpeg\r
+Content-Transfer-Encoding: base64\r
+\r
+/9j/4AAQSkZJRgABAQAAAQABAAD//gA+Q1JFQVRPUjogZ2QtanBlZyB2MS4wICh1c2luZyBJSkcg\r
+--AaB03x\r
+content-disposition: form-data; name="fileupload"; filename="bar.jpg"\r
+Content-Type: image/jpeg\r
+Content-Transfer-Encoding: base64\r
+\r
+/9j/4AAQSkZJRgABAQAAAQABAAD//gA+Q1JFQVRPUjogZ2QtanBlZyB2MS4wICh1c2luZyBJSkcg\r
+--AaB03x--\r
+EOF
+    env = Rack::MockRequest.env_for("/",
+                          "CONTENT_TYPE" => "multipart/form-data, boundary=AaB03x",
+                          "CONTENT_LENGTH" => input.size,
+                          :input => input)
+    req = Rack::Request.new(env)
+    req.params
+    env['rack.tempfiles'].size.should.equal(2)
+  end
+
   should "detect invalid multipart form data" do
     input = <<EOF
 --AaB03x\r

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