middleware.rb in pdfkit-0.5.0

- old
+ new
@@ -1,61 +1,77 @@
 class PDFKit
-  
+
   class Middleware
-    
-    def initialize(app, options = {})
-      @app = app
-      @options = options
+
+    def initialize(app, options = {}, conditions = {})
+      @app        = app
+      @options    = options
+      @conditions = conditions
     end
-        
+
     def call(env)
-      @render_pdf = false
-      set_request_to_render_as_pdf(env) if env['PATH_INFO'].match(/\.pdf$/)
-      
+      @request    = Rack::Request.new(env)
+
+      set_request_to_render_as_pdf(env) if render_as_pdf?
       status, headers, response = @app.call(env)
-      
-      request = Rack::Request.new(env)
-      if @render_pdf && headers['Content-Type'] =~ /text\/html|application\/xhtml\+xml/
-        body = response.body
-        
-        body = translate_paths(body, env)
-        
-        pdf = PDFKit.new(body, @options)
-        body = pdf.to_pdf
-        
+
+      if rendering_pdf? && headers['Content-Type'] =~ /text\/html|application\/xhtml\+xml/
+        body = response.respond_to?(:body) ? response.body : response.join
+        body = PDFKit.new(translate_paths(body, env), @options).to_pdf
+        response = [body]
+
         # Do not cache PDFs
         headers.delete('ETag')
         headers.delete('Cache-Control')
-        
-        headers["Content-Length"] = (body.respond_to?(:bytesize) ? body.bytesize : body.size).to_s
-        headers["Content-Type"] = "application/pdf"
-        
-        response = [body]
+
+        headers["Content-Length"]         = (body.respond_to?(:bytesize) ? body.bytesize : body.size).to_s
+        headers["Content-Type"]           = "application/pdf"
       end
-      
+
       [status, headers, response]
     end
-    
+
     private
-    
-      # Change relative paths to absolute
-      def translate_paths(body, env)
-        # Host with protocol
-        root = env['rack.url_scheme'] + "://" + env['HTTP_HOST'] + "/"
-        
-        body.gsub(/(href|src)=(['"])\/([^\"']*|[^"']*)['"]/,'\1=\2'+root+'\3\2')        
-      end
-    
-      def set_request_to_render_as_pdf(env)
-        @render_pdf = true
 
-        path = Pathname(env['PATH_INFO'])
-        ['PATH_INFO','REQUEST_URI'].each { |e| env[e] = path.to_s.sub(/#{path.extname}$/,'')  } if path.extname == '.pdf'
-        env['HTTP_ACCEPT'] = concat(env['HTTP_ACCEPT'], Rack::Mime.mime_type('.html'))
+    # Change relative paths to absolute
+    def translate_paths(body, env)
+      # Host with protocol
+      root = "#{env['rack.url_scheme']}://#{env['HTTP_HOST']}/"
+
+      body.gsub(/(href|src)=(['"])\/([^\"']*|[^"']*)['"]/, '\1=\2' + root + '\3\2')
+    end
+
+    def rendering_pdf?
+      @render_pdf
+    end
+
+    def render_as_pdf?
+      request_path_is_pdf = @request.path.match(%r{\.pdf$})
+
+      if request_path_is_pdf && @conditions[:only]
+        rules = [@conditions[:only]].flatten
+        rules.any? do |pattern|
+          if pattern.is_a?(Regexp)
+            @request.path =~ pattern
+          else
+            @request.path[0, pattern.length] == pattern
+          end
+        end
+      else
+        request_path_is_pdf
       end
-      
-      def concat(accepts, type)
-        (accepts || '').split(',').unshift(type).compact.join(',')
-      end
-  
+    end
+
+    def set_request_to_render_as_pdf(env)
+      @render_pdf = true
+      path = @request.path.sub(%r{\.pdf$}, '')
+      %w[PATH_INFO REQUEST_URI].each { |e| env[e] = path }
+      env['HTTP_ACCEPT'] = concat(env['HTTP_ACCEPT'], Rack::Mime.mime_type('.html'))
+      env["Rack-Middleware-PDFKit"] = "true"
+    end
+
+    def concat(accepts, type)
+      (accepts || '').split(',').unshift(type).compact.join(',')
+    end
+
   end
 end