lib/rack/chunked.rb in rack-1.2.8 vs lib/rack/chunked.rb in rack-1.3.0.beta

@@ -5,10 +5,36 @@
   # Middleware that applies chunked transfer encoding to response bodies
   # when the response does not include a Content-Length header.
   class Chunked
     include Rack::Utils
 
+    # A body wrapper that emits chunked responses
+    class Body
+      TERM = "\r\n"
+      TAIL = "0#{TERM}#{TERM}"
+
+      include Rack::Utils
+
+      def initialize(body)
+        @body = body
+      end
+
+      def each
+        term = TERM
+        @body.each do |chunk|
+          size = bytesize(chunk)
+          next if size == 0
+          yield [size.to_s(16), term, chunk, term].join
+        end
+        yield TAIL
+      end
+
+      def close
+        @body.close if @body.respond_to?(:close)
+      end
+    end
+
     def initialize(app)
       @app = app
     end
 
     def call(env)
@@ -19,31 +45,12 @@
          STATUS_WITH_NO_ENTITY_BODY.include?(status) ||
          headers['Content-Length'] ||
          headers['Transfer-Encoding']
         [status, headers, body]
       else
-        dup.chunk(status, headers, body)
+        headers.delete('Content-Length')
+        headers['Transfer-Encoding'] = 'chunked'
+        [status, headers, Body.new(body)]
       end
-    end
-
-    def chunk(status, headers, body)
-      @body = body
-      headers.delete('Content-Length')
-      headers['Transfer-Encoding'] = 'chunked'
-      [status, headers, self]
-    end
-
-    def each
-      term = "\r\n"
-      @body.each do |chunk|
-        size = bytesize(chunk)
-        next if size == 0
-        yield [size.to_s(16), term, chunk, term].join
-      end
-      yield ["0", term, "", term].join
-    end
-
-    def close
-      @body.close if @body.respond_to?(:close)
     end
   end
 end