Improve Rack::ContentLength middleware.

josevalim · josevalim · commit 2b3abc73f3e3 · 2011-05-19T09:31:18.000-04:00
1) It makes Rack::ContentLength middleware conform with the rack specification by removing the to_ary check;

2) Make Rack::ContentLength accept an extra header argument that checks for sendfile headers
diff --git a/lib/rack/content_length.rb b/lib/rack/content_length.rb
@@ -5,8 +5,9 @@ module Rack
   class ContentLength
     include Rack::Utils
 
-    def initialize(app)
+    def initialize(app, sendfile=nil)
       @app = app
+      @sendfile = sendfile
     end
 
     def call(env)
@@ -16,10 +17,14 @@ def call(env)
       if !STATUS_WITH_NO_ENTITY_BODY.include?(status.to_i) &&
          !headers['Content-Length'] &&
          !headers['Transfer-Encoding'] &&
-         body.respond_to?(:to_ary)
+         !(@sendfile && headers[@sendfile])
 
-        length = 0
-        body.each { |part| length += bytesize(part) }
+        new_body, length = [], 0
+        body.each do |part|
+          new_body << part
+          length += bytesize(part)
+        end
+        body = new_body
         headers['Content-Length'] = length.to_s
       end
 
diff --git a/test/spec_content_length.rb b/test/spec_content_length.rb
@@ -7,13 +7,13 @@
     response[1]['Content-Length'].should.equal '13'
   end
 
-  should "not set Content-Length on variable length bodies" do
+  should "set Content-Length even on variable length bodies" do
     body = lambda { "Hello World!" }
     def body.each ; yield call ; end
 
     app = lambda { |env| [200, {'Content-Type' => 'text/plain'}, body] }
     response = Rack::ContentLength.new(app).call({})
-    response[1]['Content-Length'].should.be.nil
+    response[1]['Content-Length'].should.equal '12'
   end
 
   should "not change Content-Length if it is already set" do
@@ -28,6 +28,12 @@ def body.each ; yield call ; end
     response[1]['Content-Length'].should.equal nil
   end
 
+  should "not set Content-Length on sendfile responses" do
+    app = lambda { |env| [200, {'Content-Type' => 'text/plain', 'X-Sendfile' => 'foo'}, %w(Hello World)] }
+    response = Rack::ContentLength.new(app, "X-Sendfile").call({})
+    response[1]['Content-Length'].should.equal nil
+  end
+
   should "not set Content-Length when Transfer-Encoding is chunked" do
     app = lambda { |env| [200, {'Transfer-Encoding' => 'chunked'}, []] }
     response = Rack::ContentLength.new(app).call({})
