lib/rack/deflater.rb in rack-1.5.5 vs lib/rack/deflater.rb in rack-1.6.0.beta

@@ -15,23 +15,30 @@
   # The middleware automatically detects when compression is supported
   # and allowed. For example no transformation is made when a cache
   # directive of 'no-transform' is present, or when the response status
   # code is one that doesn't allow an entity body.
   class Deflater
-    def initialize(app)
+    ##
+    # Creates Rack::Deflater middleware.
+    #
+    # [app] rack app instance
+    # [options] hash of deflater options, i.e.
+    #           'if' - a lambda enabling / disabling deflation based on returned boolean value
+    #                  e.g use Rack::Deflater, :if => lambda { |env, status, headers, body| body.length > 512 }
+    #           'include' - a list of content types that should be compressed
+    def initialize(app, options = {})
       @app = app
+
+      @condition = options[:if]
+      @compressible_types = options[:include]
     end
 
     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/ ||
-         (headers['Content-Encoding'] && headers['Content-Encoding'] !~ /\bidentity\b/)
+      unless should_deflate?(env, status, headers, body)
         return [status, headers, body]
       end
 
       request = Request.new(env)
 
@@ -56,20 +63,21 @@
         headers.delete('Content-Length')
         [status, headers, DeflateStream.new(body)]
       when "identity"
         [status, headers, body]
       when nil
-        body.close if body.respond_to?(:close)
         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]]
+        bp = Rack::BodyProxy.new([message]) { body.close if body.respond_to?(:close) }
+        [406, {"Content-Type" => "text/plain", "Content-Length" => message.length.to_s}, bp]
       end
     end
 
     class GzipStream
       def initialize(body, mtime)
         @body = body
         @mtime = mtime
+        @closed = false
       end
 
       def each(&block)
         @writer = block
         gzip  =::Zlib::GzipWriter.new(self)
@@ -77,18 +85,23 @@
         @body.each { |part|
           gzip.write(part)
           gzip.flush
         }
       ensure
-        @body.close if @body.respond_to?(:close)
         gzip.close
         @writer = nil
       end
 
       def write(data)
         @writer.call(data)
       end
+
+      def close
+        return if @closed
+        @closed = true
+        @body.close if @body.respond_to?(:close)
+      end
     end
 
     class DeflateStream
       DEFLATE_ARGS = [
         Zlib::DEFAULT_COMPRESSION,
@@ -98,19 +111,45 @@
         Zlib::DEFAULT_STRATEGY
       ]
 
       def initialize(body)
         @body = body
+        @closed = false
       end
 
       def each
-        deflater = ::Zlib::Deflate.new(*DEFLATE_ARGS)
-        @body.each { |part| yield deflater.deflate(part, Zlib::SYNC_FLUSH) }
-        yield deflater.finish
+        deflator = ::Zlib::Deflate.new(*DEFLATE_ARGS)
+        @body.each { |part| yield deflator.deflate(part, Zlib::SYNC_FLUSH) }
+        yield deflator.finish
         nil
       ensure
+        deflator.close
+      end
+
+      def close
+        return if @closed
+        @closed = true
         @body.close if @body.respond_to?(:close)
-        deflater.close
       end
+    end
+
+    private
+
+    def should_deflate?(env, status, headers, body)
+      # 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/ ||
+         (headers['Content-Encoding'] && headers['Content-Encoding'] !~ /\bidentity\b/)
+        return false
+      end
+
+      # Skip if @compressible_types are given and does not include request's content type
+      return false if @compressible_types && !(headers.has_key?('Content-Type') && @compressible_types.include?(headers['Content-Type'][/[^;]*/]))
+
+      # Skip if @condition lambda is given and evaluates to false
+      return false if @condition && !@condition.call(env, status, headers, body)
+
+      true
     end
   end
 end