Rack::Deflater streaming

2 files changed, 53 insertions, 38 deletions

diff --git a/lib/rack/deflater.rb b/lib/rack/deflater.rb
index a42b7477..14137a94 100644
--- a/lib/rack/deflater.rb
+++ b/lib/rack/deflater.rb
@@ -33,17 +33,15 @@ module Rack
 
       case encoding
       when "gzip"
+        headers['Content-Encoding'] = "gzip"
+        headers.delete('Content-Length')
         mtime = headers.key?("Last-Modified") ?
           Time.httpdate(headers["Last-Modified"]) : Time.now
-        body = self.class.gzip(body, mtime)
-        size = Rack::Utils.bytesize(body)
-        headers = headers.merge("Content-Encoding" => "gzip", "Content-Length" => size.to_s)
-        [status, headers, [body]]
+        [status, headers, GzipStream.new(body, mtime)]
       when "deflate"
-        body = self.class.deflate(body)
-        size = Rack::Utils.bytesize(body)
-        headers = headers.merge("Content-Encoding" => "deflate", "Content-Length" => size.to_s)
-        [status, headers, [body]]
+        headers['Content-Encoding'] = "deflate"
+        headers.delete('Content-Length')
+        [status, headers, DeflateStream.new(body)]
       when "identity"
         [status, headers, body]
       when nil
@@ -52,34 +50,47 @@ module Rack
       end
     end
 
-    def self.gzip(body, mtime)
-      io = StringIO.new
-      gzip = Zlib::GzipWriter.new(io)
-      gzip.mtime = mtime
+    class GzipStream
+      def initialize(body, mtime)
+        @body = body
+        @mtime = mtime
+      end
 
-      # TODO: Add streaming
-      body.each { |part| gzip << part }
+      def each(&block)
+        @writer = block
+        gzip  =::Zlib::GzipWriter.new(self)
+        gzip.mtime = @mtime
+        @body.each { |part| gzip << part }
+        @body.close if @body.respond_to?(:close)
+        gzip.close
+        @writer = nil
+      end
 
-      gzip.close
-      return io.string
+      def write(data)
+        @writer.call(data)
+      end
     end
 
-    DEFLATE_ARGS = [
-      Zlib::DEFAULT_COMPRESSION,
-      # drop the zlib header which causes both Safari and IE to choke
-     -Zlib::MAX_WBITS,
-      Zlib::DEF_MEM_LEVEL,
-      Zlib::DEFAULT_STRATEGY
-    ]
+    class DeflateStream
+      DEFLATE_ARGS = [
+        Zlib::DEFAULT_COMPRESSION,
+        # drop the zlib header which causes both Safari and IE to choke
+        -Zlib::MAX_WBITS,
+        Zlib::DEF_MEM_LEVEL,
+        Zlib::DEFAULT_STRATEGY
+      ]
 
-    # Loosely based on Mongrel's Deflate handler
-    def self.deflate(body)
-      deflater = Zlib::Deflate.new(*DEFLATE_ARGS)
-
-      # TODO: Add streaming
-      body.each { |part| deflater << part }
+      def initialize(body)
+        @body = body
+      end
 
-      return deflater.finish
+      def each
+        deflater = ::Zlib::Deflate.new(*DEFLATE_ARGS)
+        @body.each { |part| yield deflater.deflate(part) }
+        @body.close if @body.respond_to?(:close)
+        yield deflater.finish
+        nil
+      end
     end
   end
 end
diff --git a/test/spec_rack_deflater.rb b/test/spec_rack_deflater.rb
index a1cedee5..c9bb3189 100644
--- a/test/spec_rack_deflater.rb
+++ b/test/spec_rack_deflater.rb
@@ -24,10 +24,11 @@ context "Rack::Deflater" do
     response[0].should.equal(200)
     response[1].should.equal({
       "Content-Encoding" => "deflate",
-      "Content-Length" => "8",
       "Vary" => "Accept-Encoding"
     })
-    response[2].should.equal(["K\313\317OJ,\002\000"])
+    buf = ''
+    response[2].each { |part| buf << part }
+    buf.should.equal("K\313\317OJ,\002\000")
   end
 
   # TODO: This is really just a special case of the above...
@@ -37,10 +38,11 @@ context "Rack::Deflater" do
     response[0].should.equal(200)
     response[1].should.equal({
       "Content-Encoding" => "deflate",
-      "Content-Length" => "14",
       "Vary" => "Accept-Encoding"
     })
-    response[2].should.equal(["\363H\315\311\311W(\317/\312IQ\004\000"])
+    buf = ''
+    response[2].each { |part| buf << part }
+    buf.should.equal("\363H\315\311\311W(\317/\312IQ\004\000")
   end
 
   specify "should be able to gzip bodies that respond to each" do
@@ -52,11 +54,12 @@ context "Rack::Deflater" do
     response[0].should.equal(200)
     response[1].should.equal({
       "Content-Encoding" => "gzip",
-      "Content-Length" => "26",
       "Vary" => "Accept-Encoding",
     })
 
-    io = StringIO.new(response[2].join)
+    buf = ''
+    response[2].each { |part| buf << part }
+    io = StringIO.new(buf)
     gz = Zlib::GzipReader.new(io)
     gz.read.should.equal("foobar")
     gz.close
@@ -100,12 +103,13 @@ context "Rack::Deflater" do
     response[0].should.equal(200)
     response[1].should.equal({
       "Content-Encoding" => "gzip",
-      "Content-Length" => "32",
       "Vary" => "Accept-Encoding",
       "Last-Modified" => last_modified
     })
 
-    io = StringIO.new(response[2].join)
+    buf = ''
+    response[2].each { |part| buf << part }
+    io = StringIO.new(buf)
     gz = Zlib::GzipReader.new(io)
     gz.read.should.equal("Hello World!")
     gz.close