Rack::ContentLength tweaks ...

 * Adds a Content-Length header only when the body is of knownable
   length (String, Array).
 * Does nothing when Transfer-Encoding header is present in
   response.
 * Uses a Set instead of an Array for status code lookup (linear
   search through 102 elements seemed expensive).

2 files changed, 27 insertions, 20 deletions

diff --git a/lib/rack/content_length.rb b/lib/rack/content_length.rb
index 9b21d59d..2f4f2c4a 100644
--- a/lib/rack/content_length.rb
+++ b/lib/rack/content_length.rb
@@ -1,5 +1,5 @@
 module Rack
-  # Automatically sets the Content-Length header on all String bodies
+  # Sets the Content-Length header on responses with fixed-length bodies.
   class ContentLength
     def initialize(app)
       @app = app
@@ -9,21 +9,13 @@ module Rack
       status, headers, body = @app.call(env)
 
       if !Rack::Utils::STATUS_WITH_NO_ENTITY_BODY.include?(status) &&
-          !headers['Content-Length']
+         !headers['Content-Length'] &&
+         !headers['Transfer-Encoding'] &&
+         (body.respond_to?(:to_ary) || body.respond_to?(:to_str))
 
-        bytes = 0
-        string_body = true
-
-        body.each { |part|
-          unless part.kind_of?(String)
-            string_body = false
-            break
-          end
-
-          bytes += (part.respond_to?(:bytesize) ? part.bytesize : part.size)
-        }
-
-        headers['Content-Length'] = bytes.to_s if string_body
+        body = [body] if body.respond_to?(:to_str) # rack 0.4 compat
+        length = body.to_ary.inject(0) { |len, part| len + part.length }
+        headers['Content-Length'] = length.to_s
       end
 
       [status, headers, body]
diff --git a/test/spec_rack_content_length.rb b/test/spec_rack_content_length.rb
index 4a205f88..7db9345f 100644
--- a/test/spec_rack_content_length.rb
+++ b/test/spec_rack_content_length.rb
@@ -2,26 +2,41 @@ require 'rack/mock'
 require 'rack/content_length'
 
 context "Rack::ContentLength" do
-  specify "sets Content-Length if none is set" do
+  specify "sets Content-Length on String bodies if none is set" do
     app = lambda { |env| [200, {'Content-Type' => 'text/plain'}, "Hello, World!"] }
     response = Rack::ContentLength.new(app).call({})
     response[1]['Content-Length'].should.equal '13'
   end
 
-  specify "set Content-Length if steaming body" do
-    app = lambda { |env| [200, {'Content-Type' => 'text/plain'}, ["Hello, ", "World!"]] }
+  specify "sets Content-Length on Array bodies if none is set" do
+    app = lambda { |env| [200, {'Content-Type' => 'text/plain'}, ["Hello, World!"]] }
     response = Rack::ContentLength.new(app).call({})
     response[1]['Content-Length'].should.equal '13'
   end
 
+  specify "does not set Content-Length 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
+  end
+
   specify "does not change Content-Length if it is already set" do
     app = lambda { |env| [200, {'Content-Type' => 'text/plain', 'Content-Length' => '1'}, "Hello, World!"] }
     response = Rack::ContentLength.new(app).call({})
     response[1]['Content-Length'].should.equal '1'
   end
 
-  specify "does not set Content-Length if on a 304 request" do
-    app = lambda { |env| [304, {'Content-Type' => 'text/plain'}, ""] }
+  specify "does not set Content-Length on 304 responses" do
+    app = lambda { |env| [304, {'Content-Type' => 'text/plain'}, []] }
+    response = Rack::ContentLength.new(app).call({})
+    response[1]['Content-Length'].should.equal nil
+  end
+
+  specify "does not set Content-Length when Transfer-Encoding is chunked" do
+    app = lambda { |env| [200, {'Transfer-Encoding' => 'chunked'}, []] }
     response = Rack::ContentLength.new(app).call({})
     response[1]['Content-Length'].should.equal nil
   end