deflater.rb in rack-1.0.0

- old
+ new
@@ -1,87 +1,96 @@
 require "zlib"
 require "stringio"
 require "time"  # for Time.httpdate
+require 'rack/utils'
 
 module Rack
+  class Deflater
+    def initialize(app)
+      @app = app
+    end
 
-class Deflater
-  def initialize(app)
-    @app = app
-  end
+    def call(env)
+      status, headers, body = @app.call(env)
+      headers = Utils::HeaderHash.new(headers)
 
-  def call(env)
-    status, headers, body = @app.call(env)
-    headers = Utils::HeaderHash.new(headers)
+      # Skip compressing empty entity body responses and responses with
+      # no-transform set.
+      if Utils::STATUS_WITH_NO_ENTITY_BODY.include?(status) ||
+          headers['Cache-Control'].to_s =~ /\bno-transform\b/
+        return [status, headers, body]
+      end
 
-    # Skip compressing empty entity body responses and responses with
-    # no-transform set.
-    if Utils::STATUS_WITH_NO_ENTITY_BODY.include?(status) ||
-        headers['Cache-Control'].to_s =~ /\bno-transform\b/
-      return [status, headers, body]
-    end
+      request = Request.new(env)
 
-    request = Request.new(env)
+      encoding = Utils.select_best_encoding(%w(gzip deflate identity),
+                                            request.accept_encoding)
 
-    encoding = Utils.select_best_encoding(%w(gzip deflate identity),
-                                          request.accept_encoding)
+      # Set the Vary HTTP header.
+      vary = headers["Vary"].to_s.split(",").map { |v| v.strip }
+      unless vary.include?("*") || vary.include?("Accept-Encoding")
+        headers["Vary"] = vary.push("Accept-Encoding").join(",")
+      end
 
-    # Set the Vary HTTP header.
-    vary = headers["Vary"].to_s.split(",").map { |v| v.strip }
-    unless vary.include?("*") || vary.include?("Accept-Encoding")
-      headers["Vary"] = vary.push("Accept-Encoding").join(",")
+      case encoding
+      when "gzip"
+        headers['Content-Encoding'] = "gzip"
+        headers.delete('Content-Length')
+        mtime = headers.key?("Last-Modified") ?
+          Time.httpdate(headers["Last-Modified"]) : Time.now
+        [status, headers, GzipStream.new(body, mtime)]
+      when "deflate"
+        headers['Content-Encoding'] = "deflate"
+        headers.delete('Content-Length')
+        [status, headers, DeflateStream.new(body)]
+      when "identity"
+        [status, headers, body]
+      when nil
+        message = "An acceptable encoding for the requested resource #{request.fullpath} could not be found."
+        [406, {"Content-Type" => "text/plain", "Content-Length" => message.length.to_s}, [message]]
+      end
     end
 
-    case encoding
-    when "gzip"
-      mtime = if headers.key?("Last-Modified") 
-                Time.httpdate(headers["Last-Modified"]) 
-              else 
-                Time.now
-              end
-      [status,
-       headers.merge("Content-Encoding" => "gzip"),
-       self.class.gzip(body, mtime)]
-    when "deflate"
-      [status,
-       headers.merge("Content-Encoding" => "deflate"),
-       self.class.deflate(body)]
-    when "identity"
-      [status, headers, body]
-    when nil
-      message = ["An acceptable encoding for the requested resource #{request.fullpath} could not be found."]
-      [406, {"Content-Type" => "text/plain"}, message]
-    end
-  end
+    class GzipStream
+      def initialize(body, mtime)
+        @body = body
+        @mtime = mtime
+      end
 
-  def self.gzip(body, mtime)
-    io = StringIO.new
-    gzip = Zlib::GzipWriter.new(io)
-    gzip.mtime = mtime
+      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
 
-    # TODO: Add streaming
-    body.each { |part| gzip << part }
+      def write(data)
+        @writer.call(data)
+      end
+    end
 
-    gzip.close
-    return io.string
-  end
+    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
+      ]
 
-  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
-  ]
+      def initialize(body)
+        @body = body
+      end
 
-  # 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 }
-
-    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
-
 end